<commit_before>
use std::fmt;
use std::io::Read;  // Trait providing read_to_string()
use std::env;

use hyper;
use serde;
use serde_json;

// use Groups;

pub const API_VERSION: u16 = 3;

#[derive(Default, Clone, Copy, Debug)]
pub struct Pagination {
    pub page: u16,
    pub per_page: u16,
}

#[derive(Default)]
pub struct GitLab {
    scheme: String,
    domain: String,
    port: u16,
    private_token: String,
    pagination: Option<Pagination>,
    client: hyper::Client,
}

// Explicitly implement Debug trait for GitLab so we can hide the token.
impl fmt::Debug for GitLab {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "GitLab {{ scheme: {}, domain: {}, port: {}, private_token: XXXXXXXXXXXXXXXXXXXX, \
                pagination: {:?} }}",
               self.scheme,
               self.domain,
               self.port,
               self.pagination)
    }
}

impl GitLab {
    pub fn _new(scheme: &str, domain: &str, port: u16, private_token: &str) -> GitLab {
        GitLab {
            scheme: scheme.to_string(),
            domain: domain.to_string(),
            port: port,
            private_token: private_token.to_string(),
            pagination: None,
            client: match env::var("HTTP_PROXY") {
                Ok(proxy) => {
                    let proxy: Vec<&str> = proxy.trim_left_matches("http://").split(':').collect();
                    let hostname = proxy[0].to_string();
                    let port = proxy[1];

                    hyper::Client::with_http_proxy(hostname, port.parse().unwrap())
                }
                Err(_) => hyper::Client::new(),
            },
        }
    }

    pub fn new_insecure(domain: &str, private_token: &str) -> GitLab {
        warn!("Using insecure http:// protocol: Token will be sent in clear!");
        GitLab::_new("http", domain, 80, private_token)
    }

    pub fn new(domain: &str, private_token: &str) -> GitLab {
        GitLab::_new("https", domain, 443, private_token)
    }

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

    pub fn scheme(mut self, scheme: &str) -> Self {
        self.scheme = scheme.to_string();
        self
    }

    /// Build a URL used to access GitLab instance, including some parameters.
    ///
    /// # Examples
    ///
    /// Example from GitLab: https://docs.gitlab.com/ce/api/#basic-usage
    ///
    /// ```
    /// use gitlab_api::GitLab;
    ///
    /// let expected_url = "https://gitlab.example.com:\
    ///                     443/api/v3/groups?order_by=path&private_token=XXXXXXXXXXXXX";
    ///
    /// let gl = GitLab::new("gitlab.example.com", "XXXXXXXXXXXXX");
    ///
    /// assert_eq!(gl.build_url("groups?order_by=path"), expected_url);
    /// ```
    pub fn build_url(&self, query: &str) -> String {
        let params_splitter = if query.find('?').is_some() { "&" } else { "?" };
        let mut url = format!("{}://{}:{}/api/v{}/{}{}private_token={}",
                              self.scheme,
                              self.domain,
                              self.port,
                              API_VERSION,
                              query,
                              params_splitter,
                              self.private_token);

        self.pagination.as_ref().map(|pagination| {
            url.push_str(&format!("&page={}&per_page={}", pagination.page, pagination.per_page));
        });

        url
    }

    // pub fn attempt_connection(&self) -> Result<hyper::client::Response, hyper::Error> {
    //     let url = self.build_url("version");
    //     // Close connections after each GET.
    //     self.client.get(&url).header(hyper::header::Connection::close()).send()
    // }

    pub fn set_pagination(&mut self, pagination: Pagination) {
        self.pagination = Some(pagination);
    }

    pub fn get<T>(&self, query: &str) -> Result<T, serde_json::Error>
        where T: serde::Deserialize
    {
        // FIXME: Properly handle any errors. Use chain_error crate.

        let url = self.build_url(query);
        info!("url: {:?}", url);

        // Close connections after each GET.
        let mut res: hyper::client::Response = self.client
            .get(&url)
            .header(hyper::header::Connection::close())
            .send()
            .unwrap();
        info!("res.status: {:?}", res.status);
        debug!("res.headers: {:?}", res.headers);
        debug!("res.url: {:?}", res.url);

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();
        debug!("body:\n{:?}", body);

        // assert_eq!(res.status, hyper::status::StatusCode::Ok);

        serde_json::from_str(body.as_str())
    }

    pub fn version(&self) -> Result<::Version, serde_json::Error> {
        self.get::<::Version>("version")
    }

    pub fn groups(&self) -> ::groups::GroupsLister {
        ::groups::GroupsLister::new(self)
    }

    pub fn projects(&self) -> ::projects::ProjectsLister {
        ::projects::ProjectsLister::new(self)
    }

    pub fn issues(&self) -> ::issues::IssuesLister {
        ::issues::IssuesLister::new(self)
    }

    pub fn merge_requests(&self, project_id: i64) -> ::merge_requests::MergeRequestsLister {
        ::merge_requests::MergeRequestsLister::new(self, project_id)
    }

    // pub fn groups(&mut self, listing: ::groups::Listing) -> Result<Groups, serde_json::Error> {
    //     let query = listing.build_query();
    //     // self.get(&query)
    //     unimplemented!();
    // }

    // pub fn owned_groups(&mut self) -> Result<Groups, serde_json::Error> {
    //     let query = ::groups::owned_groups::Listing::new().build_query();
    //     info!("query: {:?}", query);
    //     self.get(&query)
    // }
}

// #[cfg(test)]
// mod tests {
//
// #[test]
// fn it_works() {
// }
// }
//
<commit_msg>Let `gitlab::version()` use the new error-chain error type and handling<commit_after>
use std::fmt;
use std::io::Read;  // Trait providing read_to_string()
use std::env;

use hyper;
use serde;
use serde_json;

// use Groups;

// Import the ResultExt trait, giving the chain_err() method on Result
use errors::ResultExt;

pub const API_VERSION: u16 = 3;

#[derive(Default, Clone, Copy, Debug)]
pub struct Pagination {
    pub page: u16,
    pub per_page: u16,
}

#[derive(Default)]
pub struct GitLab {
    scheme: String,
    domain: String,
    port: u16,
    private_token: String,
    pagination: Option<Pagination>,
    client: hyper::Client,
}

// Explicitly implement Debug trait for GitLab so we can hide the token.
impl fmt::Debug for GitLab {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "GitLab {{ scheme: {}, domain: {}, port: {}, private_token: XXXXXXXXXXXXXXXXXXXX, \
                pagination: {:?} }}",
               self.scheme,
               self.domain,
               self.port,
               self.pagination)
    }
}

impl GitLab {
    pub fn _new(scheme: &str, domain: &str, port: u16, private_token: &str) -> GitLab {
        GitLab {
            scheme: scheme.to_string(),
            domain: domain.to_string(),
            port: port,
            private_token: private_token.to_string(),
            pagination: None,
            client: match env::var("HTTP_PROXY") {
                Ok(proxy) => {
                    let proxy: Vec<&str> = proxy.trim_left_matches("http://").split(':').collect();
                    let hostname = proxy[0].to_string();
                    let port = proxy[1];

                    hyper::Client::with_http_proxy(hostname, port.parse().unwrap())
                }
                Err(_) => hyper::Client::new(),
            },
        }
    }

    pub fn new_insecure(domain: &str, private_token: &str) -> GitLab {
        warn!("Using insecure http:// protocol: Token will be sent in clear!");
        GitLab::_new("http", domain, 80, private_token)
    }

    pub fn new(domain: &str, private_token: &str) -> GitLab {
        GitLab::_new("https", domain, 443, private_token)
    }

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

    pub fn scheme(mut self, scheme: &str) -> Self {
        self.scheme = scheme.to_string();
        self
    }

    /// Build a URL used to access GitLab instance, including some parameters.
    ///
    /// # Examples
    ///
    /// Example from GitLab: https://docs.gitlab.com/ce/api/#basic-usage
    ///
    /// ```
    /// use gitlab_api::GitLab;
    ///
    /// let expected_url = "https://gitlab.example.com:\
    ///                     443/api/v3/groups?order_by=path&private_token=XXXXXXXXXXXXX";
    ///
    /// let gl = GitLab::new("gitlab.example.com", "XXXXXXXXXXXXX");
    ///
    /// assert_eq!(gl.build_url("groups?order_by=path"), expected_url);
    /// ```
    pub fn build_url(&self, query: &str) -> String {
        let params_splitter = if query.find('?').is_some() { "&" } else { "?" };
        let mut url = format!("{}://{}:{}/api/v{}/{}{}private_token={}",
                              self.scheme,
                              self.domain,
                              self.port,
                              API_VERSION,
                              query,
                              params_splitter,
                              self.private_token);

        self.pagination.as_ref().map(|pagination| {
            url.push_str(&format!("&page={}&per_page={}", pagination.page, pagination.per_page));
        });

        url
    }

    // pub fn attempt_connection(&self) -> Result<hyper::client::Response, hyper::Error> {
    //     let url = self.build_url("version");
    //     // Close connections after each GET.
    //     self.client.get(&url).header(hyper::header::Connection::close()).send()
    // }

    pub fn set_pagination(&mut self, pagination: Pagination) {
        self.pagination = Some(pagination);
    }

    pub fn get<T>(&self, query: &str) -> Result<T, serde_json::Error>
        where T: serde::Deserialize
    {
        // FIXME: Properly handle any errors. Use chain_error crate.

        let url = self.build_url(query);
        info!("url: {:?}", url);

        // Close connections after each GET.
        let mut res: hyper::client::Response = self.client
            .get(&url)
            .header(hyper::header::Connection::close())
            .send()
            .unwrap();
        info!("res.status: {:?}", res.status);
        debug!("res.headers: {:?}", res.headers);
        debug!("res.url: {:?}", res.url);

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();
        debug!("body:\n{:?}", body);

        // assert_eq!(res.status, hyper::status::StatusCode::Ok);

        serde_json::from_str(body.as_str())
    }

    pub fn new_get<T>(&self, query: &str) -> ::errors::Result<T>
        where T: serde::Deserialize
    {
        let url = self.build_url(query);
        info!("url: {:?}", url);

        // Close connections after each GET.
        let mut res: hyper::client::Response = self.client
            .get(&url)
            .header(hyper::header::Connection::close())
            .send()
            .chain_err(|| format!("cannot send request {} to {:?}", query, self))?;
        info!("res.status: {:?}", res.status);
        debug!("res.headers: {:?}", res.headers);
        debug!("res.url: {:?}", res.url);

        let mut body = String::new();
        res.read_to_string(&mut body).chain_err(|| "cannot read response body")?;
        debug!("body:\n{:?}", body);

        // assert_eq!(res.status, hyper::status::StatusCode::Ok);
        if res.status != hyper::status::StatusCode::Ok {
            bail!(format!("status code ({}) not Ok()", res.status));
        }

        serde_json::from_str(body.as_str()).chain_err(|| format!("cannot build Rust struct from JSON data: {}", body))
    }

    pub fn version(&self) -> ::errors::Result<::Version> {
        self.new_get("version").chain_err(|| "cannot query 'version'")
    }

    pub fn groups(&self) -> ::groups::GroupsLister {
        ::groups::GroupsLister::new(self)
    }

    pub fn projects(&self) -> ::projects::ProjectsLister {
        ::projects::ProjectsLister::new(self)
    }

    pub fn issues(&self) -> ::issues::IssuesLister {
        ::issues::IssuesLister::new(self)
    }

    pub fn merge_requests(&self, project_id: i64) -> ::merge_requests::MergeRequestsLister {
        ::merge_requests::MergeRequestsLister::new(self, project_id)
    }

    // pub fn groups(&mut self, listing: ::groups::Listing) -> Result<Groups, serde_json::Error> {
    //     let query = listing.build_query();
    //     // self.get(&query)
    //     unimplemented!();
    // }

    // pub fn owned_groups(&mut self) -> Result<Groups, serde_json::Error> {
    //     let query = ::groups::owned_groups::Listing::new().build_query();
    //     info!("query: {:?}", query);
    //     self.get(&query)
    // }
}

// #[cfg(test)]
// mod tests {
//
// #[test]
// fn it_works() {
// }
// }
//
<|endoftext|>


<commit_before>use std::vec::IntoIter;

const TOPIC_PATH_DELIMITER: char = '/';

use self::Topic::{
    Normal,
    System,
    Blank,
    SingleWildcard,
    MultiWildcard
};

/// FIXME: replace String with &str
#[derive(Debug, Clone, PartialEq)]
pub enum Topic {
    Normal(String),
    System(String), // $SYS = Topic::System("$SYS")
    Blank,
    SingleWildcard, // +
    MultiWildcard,  // #
}

impl Into<String> for Topic {
    fn into(self) -> String {
        match self {
            Normal(s) | System(s) => s,
            Blank => "".to_string(),
            SingleWildcard => "+".to_string(),
            MultiWildcard => "#".to_string()
        }
    }
}

impl Topic {
    pub fn fit(&self, other: &Topic) -> bool {
        match *self {
            Normal(ref str) => {
                match *other {
                    Normal(ref s) => str == s,
                    SingleWildcard | MultiWildcard => true,
                    _ => false
                }
            },
            System(ref str) => {
                match *other {
                    System(ref s) => str == s,
                    _ => false
                }
            },
            Blank => {
                match *other {
                    Blank | SingleWildcard | MultiWildcard => true,
                    _ => false
                }
            },
            SingleWildcard => {
                match *other {
                    System(_) => false,
                    _ => true
                }
            },
            MultiWildcard => {
                match *other {
                    System(_) => false,
                    _ => true
                }
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct TopicPath {
    pub path: String,
    // Should be false for Topic Name
    pub wildcards: bool,
    topics: Vec<Topic>
}

impl TopicPath {
    pub fn path(&self) -> String {
        self.path.clone()
    }

    pub fn get(&self, index: usize) -> Option<&Topic> {
        self.topics.get(index)
    }

    pub fn get_mut(&mut self, index: usize) -> Option<&mut Topic> {
        self.topics.get_mut(index)
    }

    pub fn len(&self) -> usize {
        self.topics.len()
    }

    pub fn into_iter(self) -> IntoIter<Topic> {
        self.topics.into_iter()
    }

    pub fn is_final(&self, index: usize) -> bool {
        let len = self.topics.len();
        len == 0 || len-1 == index
    }

    pub fn is_multi(&self, index: usize) -> bool {
        match self.topics.get(index) {
            Some(topic) => *topic == Topic::MultiWildcard,
            None => false
        }
    }
}

impl<'a> From<&'a str> for TopicPath {
    fn from(str: &'a str) -> TopicPath {
        Self::from(String::from(str))
    }
}

impl From<String> for TopicPath {
    fn from(path: String) -> TopicPath {
        let topics: Vec<Topic> = path.split(TOPIC_PATH_DELIMITER).map( |topic| {
            match topic {
                "+" => Topic::SingleWildcard,
                "#" => Topic::MultiWildcard,
                "" => Topic::Blank,
                _ => {
                    if topic.chars().nth(0) == Some('$') {
                        Topic::System(String::from(topic))
                    } else {
                        // FIXME: validate the topic
                        Topic::Normal(String::from(topic))
                    }
                }
            }
        }).collect();

        // check for wildcards
        let wildcards = topics.iter().any(|topic| {
            match *topic {
                Topic::SingleWildcard | Topic::MultiWildcard => true,
                _ => false
            }
        });

        TopicPath {
            path: path,
            topics: topics,
            wildcards: wildcards
        }
    }
}

impl Into<String> for TopicPath {
    fn into(self) -> String {
        self.path
    }
}

pub trait ToTopicPath {
    fn to_topic_path(&self) -> TopicPath;
}

impl ToTopicPath for TopicPath {
    fn to_topic_path(&self) -> TopicPath {
        self.clone()
    }
}

impl ToTopicPath for String {
    fn to_topic_path(&self) -> TopicPath {
        TopicPath::from(self.clone())
    }
}

impl<'a> ToTopicPath for &'a str {
    fn to_topic_path(&self) -> TopicPath {
        TopicPath::from(*self)
    }
}

#[cfg(test)]
mod test {
    use super::{TopicPath, Topic};

    #[test]
    fn topic_path_test() {
        let path = "/$SYS/test/+/#";
        let topic = TopicPath::from(path);
        let mut iter = topic.into_iter();
        assert_eq!(iter.next().unwrap(), Topic::Blank);
        assert_eq!(iter.next().unwrap(), Topic::System("$SYS".to_string()));
        assert_eq!(iter.next().unwrap(), Topic::Normal("test".to_string()));
        assert_eq!(iter.next().unwrap(), Topic::SingleWildcard);
        assert_eq!(iter.next().unwrap(), Topic::MultiWildcard);
    }

    #[test]
    fn wildcards_test() {
        let topic = TopicPath::from("/a/b/c");
        assert!(!topic.wildcards);
        let topic = TopicPath::from("/a/+/c");
        assert!(topic.wildcards);
        let topic = TopicPath::from("/a/b/#");
        assert!(topic.wildcards);
    }
}
<commit_msg>Use TopicPath::from_str instead TopicPath::from<commit_after>use std::vec::IntoIter;
use {Error, Result};

const TOPIC_PATH_DELIMITER: char = '/';

use self::Topic::{
    Normal,
    System,
    Blank,
    SingleWildcard,
    MultiWildcard
};

/// FIXME: replace String with &str
#[derive(Debug, Clone, PartialEq)]
pub enum Topic {
    Normal(String),
    System(String), // $SYS = Topic::System("$SYS")
    Blank,
    SingleWildcard, // +
    MultiWildcard,  // #
}

impl Into<String> for Topic {
    fn into(self) -> String {
        match self {
            Normal(s) | System(s) => s,
            Blank => "".to_string(),
            SingleWildcard => "+".to_string(),
            MultiWildcard => "#".to_string()
        }
    }
}

impl Topic {
    pub fn fit(&self, other: &Topic) -> bool {
        match *self {
            Normal(ref str) => {
                match *other {
                    Normal(ref s) => str == s,
                    SingleWildcard | MultiWildcard => true,
                    _ => false
                }
            },
            System(ref str) => {
                match *other {
                    System(ref s) => str == s,
                    _ => false
                }
            },
            Blank => {
                match *other {
                    Blank | SingleWildcard | MultiWildcard => true,
                    _ => false
                }
            },
            SingleWildcard => {
                match *other {
                    System(_) => false,
                    _ => true
                }
            },
            MultiWildcard => {
                match *other {
                    System(_) => false,
                    _ => true
                }
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct TopicPath {
    pub path: String,
    // Should be false for Topic Name
    pub wildcards: bool,
    topics: Vec<Topic>
}

impl TopicPath {
    pub fn path(&self) -> String {
        self.path.clone()
    }

    pub fn get(&self, index: usize) -> Option<&Topic> {
        self.topics.get(index)
    }

    pub fn get_mut(&mut self, index: usize) -> Option<&mut Topic> {
        self.topics.get_mut(index)
    }

    pub fn len(&self) -> usize {
        self.topics.len()
    }

    pub fn into_iter(self) -> IntoIter<Topic> {
        self.topics.into_iter()
    }

    pub fn is_final(&self, index: usize) -> bool {
        let len = self.topics.len();
        len == 0 || len-1 == index
    }

    pub fn is_multi(&self, index: usize) -> bool {
        match self.topics.get(index) {
            Some(topic) => *topic == Topic::MultiWildcard,
            None => false
        }
    }

    pub fn from_str<T: AsRef<str>>(path: T) -> Result<TopicPath> {
        let topics: Vec<Topic> = path.as_ref().split(TOPIC_PATH_DELIMITER).map( |topic| {
            match topic {
                "+" => Topic::SingleWildcard,
                "#" => Topic::MultiWildcard,
                "" => Topic::Blank,
                _ => {
                    if topic.chars().nth(0) == Some('$') {
                        Topic::System(String::from(topic))
                    } else {
                        // FIXME: validate the topic
                        Topic::Normal(String::from(topic))
                    }
                }
            }
        }).collect();

        // check for wildcards
        let wildcards = topics.iter().any(|topic| {
            match *topic {
                Topic::SingleWildcard | Topic::MultiWildcard => true,
                _ => false
            }
        });

        Ok(TopicPath {
            path: String::from(path.as_ref()),
            topics: topics,
            wildcards: wildcards
        })
    }
}

impl<'a> From<&'a str> for TopicPath {
    fn from(str: &'a str) -> TopicPath {
        Self::from_str(str).unwrap()
    }
}

impl From<String> for TopicPath {
    fn from(path: String) -> TopicPath {
        Self::from_str(path).unwrap()
    }
}

impl Into<String> for TopicPath {
    fn into(self) -> String {
        self.path
    }
}

pub trait ToTopicPath {
    fn to_topic_path(&self) -> Result<TopicPath>;

    fn to_topic_name(&self) -> Result<TopicPath> {
        let topic_name = try!(self.to_topic_path());
        match topic_name.wildcards {
            false => Ok(topic_name),
            true => Err(Error::TopicNameMustNotContainWildcard)
        }
    }
}

impl ToTopicPath for TopicPath {
    fn to_topic_path(&self) -> Result<TopicPath> {
        Ok(self.clone())
    }
}

impl ToTopicPath for String {
    fn to_topic_path(&self) -> Result<TopicPath> {
        TopicPath::from_str(self.clone())
    }
}

impl<'a> ToTopicPath for &'a str {
    fn to_topic_path(&self) -> Result<TopicPath> {
        TopicPath::from_str(*self)
    }
}

#[cfg(test)]
mod test {
    use super::{TopicPath, Topic};

    #[test]
    fn topic_path_test() {
        let path = "/$SYS/test/+/#";
        let topic = TopicPath::from(path);
        let mut iter = topic.into_iter();
        assert_eq!(iter.next().unwrap(), Topic::Blank);
        assert_eq!(iter.next().unwrap(), Topic::System("$SYS".to_string()));
        assert_eq!(iter.next().unwrap(), Topic::Normal("test".to_string()));
        assert_eq!(iter.next().unwrap(), Topic::SingleWildcard);
        assert_eq!(iter.next().unwrap(), Topic::MultiWildcard);
    }

    #[test]
    fn wildcards_test() {
        let topic = TopicPath::from("/a/b/c");
        assert!(!topic.wildcards);
        let topic = TopicPath::from("/a/+/c");
        assert!(topic.wildcards);
        let topic = TopicPath::from("/a/b/#");
        assert!(topic.wildcards);
    }
}
<|endoftext|>


<commit_before>
use size_hint;

use std::cmp::Ordering;
use std::collections::BinaryHeap;

macro_rules! clone_fields {
    ($name:ident, $base:expr, $($field:ident),+) => (
        $name {
            $(
                $field : $base . $field .clone()
            ),*
        }
    );
}

/// A non-empty sequence
///
/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
/// first items (which are guaranteed to exist).
///
/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the `BinaryHeap` used in
/// `KMerge` into a min-heap.
struct NonEmpty<I>
    where I: Iterator
{
    head: I::Item,
    tail: I,
}

impl<I> NonEmpty<I>
    where I: Iterator
{
    /// Constructs a `NonEmpty` from an `Iterator`. Returns `None` if the `Iterator` is empty.
    fn new(mut it: I) -> Option<NonEmpty<I>> {
        let head = it.next();
        head.map(|h| {
            NonEmpty {
                head: h,
                tail: it,
            }
        })
    }

    /// Returns the next item in the sequence. If more items remain, the remainder of the sequence
    /// is returned as well, otherwise `None`.
    fn next(self) -> (I::Item, Option<Self>) {
        (self.head, NonEmpty::new(self.tail))
    }

    /// Hints at the size of the sequence, same as the `Iterator` method.
    fn size_hint(&self) -> (usize, Option<usize>) {
        size_hint::add_scalar(self.tail.size_hint(), 1)
    }
}

impl<I> Clone for NonEmpty<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> Self {
        clone_fields!(NonEmpty, self, head, tail)
    }
}

impl<I> PartialEq for NonEmpty<I>
    where I: Iterator,
          I::Item: PartialEq
{
    fn eq(&self, other: &NonEmpty<I>) -> bool {
        self.head.eq(&other.head)
    }
}

impl<I> Eq for NonEmpty<I>
    where I: Iterator,
          I::Item: Eq
{}

impl<I> PartialOrd for NonEmpty<I>
    where I: Iterator,
          I::Item: PartialOrd
{
    fn partial_cmp(&self, other: &NonEmpty<I>) -> Option<Ordering> {
        other.head.partial_cmp(&self.head)
    }

    fn lt(&self, other: &NonEmpty<I>) -> bool {
        other.head.lt(&self.head)
    }

    fn le(&self, other: &NonEmpty<I>) -> bool {
        other.head.le(&self.head)
    }

    fn gt(&self, other: &NonEmpty<I>) -> bool {
        other.head.gt(&self.head)
    }

    fn ge(&self, other: &NonEmpty<I>) -> bool {
        other.head.ge(&self.head)
    }
}

impl<I> Ord for NonEmpty<I>
    where I: Iterator,
          I::Item: Ord
{
    fn cmp(&self, other: &Self) -> Ordering {
        other.head.cmp(&self.head)
    }
}

/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
/// If all base iterators are sorted (ascending), the result is sorted.
///
/// Iterator element type is `I::Item`.
///
/// See [*.kmerge()*](trait.Itertools.html#method.kmerge) for more information.
pub struct KMerge<I>
    where I: Iterator
{
    heap: BinaryHeap<NonEmpty<I>>,
}

/// Create a `KMerge` iterator.
pub fn kmerge_new<I>(it: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
    where I: Iterator,
          I::Item: IntoIterator,
          <<I as Iterator>::Item as IntoIterator>::Item: Ord
{
    KMerge { heap: it.filter_map(|it| NonEmpty::new(it.into_iter())).collect() }
}

impl<I> Clone for KMerge<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> KMerge<I> {
        clone_fields!(KMerge, self, heap)
    }
}

impl<I> Iterator for KMerge<I>
    where I: Iterator,
          I::Item: Ord
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        self.heap.pop().map(|p| {
            let (h, t) = p.next();
            t.map(|t| self.heap.push(t));
            h
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        if !self.heap.is_empty() {
            let mut it = self.heap.iter();
            let sh0 = it.next().unwrap().size_hint();
            it.fold(sh0, |acc, it| size_hint::add(acc, it.size_hint()))
        } else {
            (0, Some(0))
        }
    }
}

<commit_msg>Use a simple heap / sift down implementation to optimize kmerge<commit_after>
use size_hint;

use std::cmp::Ordering;
use std::mem::replace;

macro_rules! clone_fields {
    ($name:ident, $base:expr, $($field:ident),+) => (
        $name {
            $(
                $field : $base . $field .clone()
            ),*
        }
    );
}

/// Head element and Tail iterator pair
///
/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
/// first items (which are guaranteed to exist).
///
/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
/// `KMerge` into a min-heap.
struct HeadTail<I>
    where I: Iterator
{
    head: I::Item,
    tail: I,
}

impl<I> HeadTail<I>
    where I: Iterator
{
    /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
    fn new(mut it: I) -> Option<HeadTail<I>> {
        let head = it.next();
        head.map(|h| {
            HeadTail {
                head: h,
                tail: it,
            }
        })
    }

    /// Hints at the size of the sequence, same as the `Iterator` method.
    fn size_hint(&self) -> (usize, Option<usize>) {
        size_hint::add_scalar(self.tail.size_hint(), 1)
    }
}

impl<I> Clone for HeadTail<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> Self {
        clone_fields!(HeadTail, self, head, tail)
    }
}

impl<I> PartialEq for HeadTail<I>
    where I: Iterator,
          I::Item: PartialEq
{
    fn eq(&self, other: &HeadTail<I>) -> bool {
        self.head.eq(&other.head)
    }
}

impl<I> Eq for HeadTail<I>
    where I: Iterator,
          I::Item: Eq
{}

impl<I> PartialOrd for HeadTail<I>
    where I: Iterator,
          I::Item: PartialOrd
{
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.head.partial_cmp(&self.head)
    }

    fn lt(&self, other: &Self) -> bool {
        other.head.lt(&self.head)
    }

    fn le(&self, other: &Self) -> bool {
        other.head.le(&self.head)
    }

    fn gt(&self, other: &Self) -> bool {
        other.head.gt(&self.head)
    }

    fn ge(&self, other: &Self) -> bool {
        other.head.ge(&self.head)
    }
}

impl<I> Ord for HeadTail<I>
    where I: Iterator,
          I::Item: Ord
{
    fn cmp(&self, other: &Self) -> Ordering {
        other.head.cmp(&self.head)
    }
}

/// Make `data` a heap (max-heap w.r.t T's Ord).
fn heapify<T: Ord>(data: &mut [T]) {
    for i in 0..data.len() / 2 {
        sift_down(data, i);
    }
}

/// Sift down element at `index` (`heap` is a max-heap wrt T's Ord).
fn sift_down<T: Ord>(heap: &mut [T], index: usize) {
    debug_assert!(index <= heap.len());
    let mut pos = index;
    let mut child = 2 * pos + 1;
    // the `pos` conditional is to avoid a bounds check
    while pos < heap.len() && child < heap.len() {
        let right = child + 1;

        // pick the bigger of the two children
        if right < heap.len() && heap[child] < heap[right] {
            child = right;
        }

        // sift down is done if we are already in order
        if heap[pos] > heap[child] {
            return;
        }
        heap.swap(pos, child);
        pos = child;
        child = 2 * pos + 1;
    }
}

/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
/// If all base iterators are sorted (ascending), the result is sorted.
///
/// Iterator element type is `I::Item`.
///
/// See [*.kmerge()*](trait.Itertools.html#method.kmerge) for more information.
pub struct KMerge<I>
    where I: Iterator
{
    heap: Vec<HeadTail<I>>,
}

/// Create a `KMerge` iterator.
pub fn kmerge_new<I>(iter: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
    where I: Iterator,
          I::Item: IntoIterator,
          <<I as Iterator>::Item as IntoIterator>::Item: Ord
{
    let (lower, _) = iter.size_hint();
    let mut heap = Vec::with_capacity(lower);
    heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
    heapify(&mut heap);
    KMerge { heap: heap }
}

impl<I> Clone for KMerge<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> KMerge<I> {
        clone_fields!(KMerge, self, heap)
    }
}

impl<I> Iterator for KMerge<I>
    where I: Iterator,
          I::Item: Ord
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.heap.is_empty() {
            return None;
        }
        let result = if let Some(next_0) = self.heap[0].tail.next() {
            replace(&mut self.heap[0].head, next_0)
        } else {
            self.heap.swap_remove(0).head
        };
        sift_down(&mut self.heap, 0);
        Some(result)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        if !self.heap.is_empty() {
            let mut it = self.heap.iter();
            let sh0 = it.next().unwrap().size_hint();
            it.fold(sh0, |acc, it| size_hint::add(acc, it.size_hint()))
        } else {
            (0, Some(0))
        }
    }
}

<|endoftext|>


<commit_before>// Copyright 2014 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//! A one-dimensional length, tagged with its units.

use scale_factor::ScaleFactor;
use num::Zero;

use heapsize::HeapSizeOf;
use num_traits::NumCast;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::cmp::Ordering;
use std::ops::{Add, Sub, Mul, Div, Neg};
use std::ops::{AddAssign, SubAssign};
use std::marker::PhantomData;
use std::fmt;

/// A one-dimensional distance, with value represented by `T` and unit of measurement `Unit`.
///
/// `T` can be any numeric type, for example a primitive type like u64 or f32.
///
/// `Unit` is not used in the representation of a Length value. It is used only at compile time
/// to ensure that a Length stored with one unit is converted explicitly before being used in an
/// expression that requires a different unit.  It may be a type without values, such as an empty
/// enum.
///
/// You can multiply a Length by a `scale_factor::ScaleFactor` to convert it from one unit to
/// another. See the `ScaleFactor` docs for an example.
// Uncomment the derive, and remove the macro call, once heapsize gets
// PhantomData<T> support.
#[derive(RustcDecodable, RustcEncodable)]
pub struct Length<T, Unit>(pub T, PhantomData<Unit>);

impl<T: Clone, Unit> Clone for Length<T, Unit> {
    fn clone(&self) -> Self {
        Length(self.0.clone(), PhantomData)
    }
}

impl<T: Copy, Unit> Copy for Length<T, Unit> {}

impl<Unit, T: HeapSizeOf> HeapSizeOf for Length<T, Unit> {
    fn heap_size_of_children(&self) -> usize {
        self.0.heap_size_of_children()
    }
}

impl<Unit, T> Deserialize for Length<T, Unit> where T: Deserialize {
    fn deserialize<D>(deserializer: &mut D) -> Result<Length<T, Unit>,D::Error>
                      where D: Deserializer {
        Ok(Length(try!(Deserialize::deserialize(deserializer)), PhantomData))
    }
}

impl<T, Unit> Serialize for Length<T, Unit> where T: Serialize {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(),S::Error> where S: Serializer {
        self.0.serialize(serializer)
    }
}

impl<T, Unit> Length<T, Unit> {
    pub fn new(x: T) -> Length<T, Unit> {
        Length(x, PhantomData)
    }
}

impl<Unit, T: Clone> Length<T, Unit> {
    pub fn get(&self) -> T {
        self.0.clone()
    }
}

impl<T: fmt::Debug + Clone, U> fmt::Debug for Length<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.get().fmt(f)
    }
}

impl<T: fmt::Display + Clone, U> fmt::Display for Length<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.get().fmt(f)
    }
}

// length + length
impl<U, T: Clone + Add<T, Output=T>> Add for Length<T, U> {
    type Output = Length<T, U>;
    fn add(self, other: Length<T, U>) -> Length<T, U> {
        Length::new(self.get() + other.get())
    }
}

// length += length
impl<U, T: Clone + AddAssign<T>> AddAssign for Length<T, U> {
    fn add_assign(&mut self, other: Length<T, U>) {
        self.0 += other.get();
    }
}

// length - length
impl<U, T: Clone + Sub<T, Output=T>> Sub<Length<T, U>> for Length<T, U> {
    type Output = Length<T, U>;
    fn sub(self, other: Length<T, U>) -> <Self as Sub>::Output {
        Length::new(self.get() - other.get())
    }
}

// length -= length
impl<U, T: Clone + SubAssign<T>> SubAssign for Length<T, U> {
    fn sub_assign(&mut self, other: Length<T, U>) {
        self.0 -= other.get();
    }
}

// length / length
impl<Src, Dst, T: Clone + Div<T, Output=T>> Div<Length<T, Src>> for Length<T, Dst> {
    type Output = ScaleFactor<T, Src, Dst>;
    #[inline]
    fn div(self, other: Length<T, Src>) -> ScaleFactor<T, Src, Dst> {
        ScaleFactor::new(self.get() / other.get())
    }
}

// length * scaleFactor
impl<Src, Dst, T: Clone + Mul<T, Output=T>> Mul<ScaleFactor<T, Src, Dst>> for Length<T, Src> {
    type Output = Length<T, Dst>;
    #[inline]
    fn mul(self, scale: ScaleFactor<T, Src, Dst>) -> Length<T, Dst> {
        Length::new(self.get() * scale.get())
    }
}

// length / scaleFactor
impl<Src, Dst, T: Clone + Div<T, Output=T>> Div<ScaleFactor<T, Src, Dst>> for Length<T, Dst> {
    type Output = Length<T, Src>;
    #[inline]
    fn div(self, scale: ScaleFactor<T, Src, Dst>) -> Length<T, Src> {
        Length::new(self.get() / scale.get())
    }
}

// -length
impl <U, T:Clone + Neg<Output=T>> Neg for Length<T, U> {
    type Output = Length<T, U>;
    #[inline]
    fn neg(self) -> Length<T, U> {
        Length::new(-self.get())
    }
}

impl<Unit, T0: NumCast + Clone> Length<T0, Unit> {
    /// Cast from one numeric representation to another, preserving the units.
    pub fn cast<T1: NumCast + Clone>(&self) -> Option<Length<T1, Unit>> {
        NumCast::from(self.get()).map(Length::new)
    }
}

impl<Unit, T: Clone + PartialEq> PartialEq for Length<T, Unit> {
    fn eq(&self, other: &Length<T, Unit>) -> bool { self.get().eq(&other.get()) }
}

impl<Unit, T: Clone + PartialOrd> PartialOrd for Length<T, Unit> {
    fn partial_cmp(&self, other: &Length<T, Unit>) -> Option<Ordering> {
        self.get().partial_cmp(&other.get())
    }
}

impl<Unit, T: Clone + Eq> Eq for Length<T, Unit> {}

impl<Unit, T: Clone + Ord> Ord for Length<T, Unit> {
    fn cmp(&self, other: &Length<T, Unit>) -> Ordering { self.get().cmp(&other.get()) }
}

impl<Unit, T: Zero> Zero for Length<T, Unit> {
    fn zero() -> Length<T, Unit> {
        Length::new(Zero::zero())
    }
}

#[cfg(test)]
mod tests {
    use super::Length;
    use scale_factor::ScaleFactor;

    enum Inch {}
    enum Mm {}

    #[test]
    fn test_length() {
        let mm_per_inch: ScaleFactor<f32, Inch, Mm> = ScaleFactor::new(25.4);

        let one_foot: Length<f32, Inch> = Length::new(12.0);
        let two_feet = one_foot.clone() + one_foot.clone();
        let zero_feet = one_foot.clone() - one_foot.clone();

        assert_eq!(one_foot.get(), 12.0);
        assert_eq!(two_feet.get(), 24.0);
        assert_eq!(zero_feet.get(), 0.0);

        assert!(one_foot == one_foot);
        assert!(two_feet != one_foot);

        assert!(zero_feet <  one_foot);
        assert!(zero_feet <= one_foot);
        assert!(two_feet  >  one_foot);
        assert!(two_feet  >= one_foot);

        assert!(  two_feet <= two_feet);
        assert!(  two_feet >= two_feet);
        assert!(!(two_feet >  two_feet));
        assert!(!(two_feet <  two_feet));

        let one_foot_in_mm: Length<f32, Mm> = one_foot * mm_per_inch;

        assert_eq!(one_foot_in_mm, Length::new(304.8));
        assert_eq!(one_foot_in_mm / one_foot, mm_per_inch);

        let back_to_inches: Length<f32, Inch> = one_foot_in_mm / mm_per_inch;
        assert_eq!(one_foot, back_to_inches);

        let int_foot: Length<isize, Inch> = one_foot.cast().unwrap();
        assert_eq!(int_foot.get(), 12);

        let negative_one_foot = -one_foot;
        assert_eq!(negative_one_foot.get(), -12.0);

        let negative_two_feet = -two_feet;
        assert_eq!(negative_two_feet.get(), -24.0);

        let zero_feet: Length<f32, Inch> = Length::new(0.0);
        let negative_zero_feet = -zero_feet;
        assert_eq!(negative_zero_feet.get(), 0.0);
    }

    #[test]
    fn test_addassign() {
        let one_cm: Length<f32, Mm> = Length::new(10.0);
        let mut measurement: Length<f32, Mm> = Length::new(5.0);

        measurement += one_cm;

        assert_eq!(measurement.get(), 15.0);
    }

    #[test]
    fn test_subassign() {
        let one_cm: Length<f32, Mm> = Length::new(10.0);
        let mut measurement: Length<f32, Mm> = Length::new(5.0);

        measurement -= one_cm;

        assert_eq!(measurement.get(), -5.0);
    }
}
<commit_msg>Add support for Saturating add and subtract.<commit_after>// Copyright 2014 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//! A one-dimensional length, tagged with its units.

use scale_factor::ScaleFactor;
use num::Zero;

use heapsize::HeapSizeOf;
use num_traits::{NumCast, Saturating};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::cmp::Ordering;
use std::ops::{Add, Sub, Mul, Div, Neg};
use std::ops::{AddAssign, SubAssign};
use std::marker::PhantomData;
use std::fmt;

/// A one-dimensional distance, with value represented by `T` and unit of measurement `Unit`.
///
/// `T` can be any numeric type, for example a primitive type like u64 or f32.
///
/// `Unit` is not used in the representation of a Length value. It is used only at compile time
/// to ensure that a Length stored with one unit is converted explicitly before being used in an
/// expression that requires a different unit.  It may be a type without values, such as an empty
/// enum.
///
/// You can multiply a Length by a `scale_factor::ScaleFactor` to convert it from one unit to
/// another. See the `ScaleFactor` docs for an example.
// Uncomment the derive, and remove the macro call, once heapsize gets
// PhantomData<T> support.
#[derive(RustcDecodable, RustcEncodable)]
pub struct Length<T, Unit>(pub T, PhantomData<Unit>);

impl<T: Clone, Unit> Clone for Length<T, Unit> {
    fn clone(&self) -> Self {
        Length(self.0.clone(), PhantomData)
    }
}

impl<T: Copy, Unit> Copy for Length<T, Unit> {}

impl<Unit, T: HeapSizeOf> HeapSizeOf for Length<T, Unit> {
    fn heap_size_of_children(&self) -> usize {
        self.0.heap_size_of_children()
    }
}

impl<Unit, T> Deserialize for Length<T, Unit> where T: Deserialize {
    fn deserialize<D>(deserializer: &mut D) -> Result<Length<T, Unit>,D::Error>
                      where D: Deserializer {
        Ok(Length(try!(Deserialize::deserialize(deserializer)), PhantomData))
    }
}

impl<T, Unit> Serialize for Length<T, Unit> where T: Serialize {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(),S::Error> where S: Serializer {
        self.0.serialize(serializer)
    }
}

impl<T, Unit> Length<T, Unit> {
    pub fn new(x: T) -> Length<T, Unit> {
        Length(x, PhantomData)
    }
}

impl<Unit, T: Clone> Length<T, Unit> {
    pub fn get(&self) -> T {
        self.0.clone()
    }
}

impl<T: fmt::Debug + Clone, U> fmt::Debug for Length<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.get().fmt(f)
    }
}

impl<T: fmt::Display + Clone, U> fmt::Display for Length<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.get().fmt(f)
    }
}

// length + length
impl<U, T: Clone + Add<T, Output=T>> Add for Length<T, U> {
    type Output = Length<T, U>;
    fn add(self, other: Length<T, U>) -> Length<T, U> {
        Length::new(self.get() + other.get())
    }
}

// length += length
impl<U, T: Clone + AddAssign<T>> AddAssign for Length<T, U> {
    fn add_assign(&mut self, other: Length<T, U>) {
        self.0 += other.get();
    }
}

// length - length
impl<U, T: Clone + Sub<T, Output=T>> Sub<Length<T, U>> for Length<T, U> {
    type Output = Length<T, U>;
    fn sub(self, other: Length<T, U>) -> <Self as Sub>::Output {
        Length::new(self.get() - other.get())
    }
}

// length -= length
impl<U, T: Clone + SubAssign<T>> SubAssign for Length<T, U> {
    fn sub_assign(&mut self, other: Length<T, U>) {
        self.0 -= other.get();
    }
}

// Saturating length + length and length - length.
impl<U, T: Clone + Saturating> Saturating for Length<T, U> {
    // type Output = Length<T, U>;
    fn saturating_add(self, other: Length<T, U>) -> Length<T, U> {
        Length::new(self.get().saturating_add(other.get()))
    }

    fn saturating_sub(self, other: Length<T, U>) -> Length<T, U> {
        Length::new(self.get().saturating_sub(other.get()))
    }
}

// length / length
impl<Src, Dst, T: Clone + Div<T, Output=T>> Div<Length<T, Src>> for Length<T, Dst> {
    type Output = ScaleFactor<T, Src, Dst>;
    #[inline]
    fn div(self, other: Length<T, Src>) -> ScaleFactor<T, Src, Dst> {
        ScaleFactor::new(self.get() / other.get())
    }
}

// length * scaleFactor
impl<Src, Dst, T: Clone + Mul<T, Output=T>> Mul<ScaleFactor<T, Src, Dst>> for Length<T, Src> {
    type Output = Length<T, Dst>;
    #[inline]
    fn mul(self, scale: ScaleFactor<T, Src, Dst>) -> Length<T, Dst> {
        Length::new(self.get() * scale.get())
    }
}

// length / scaleFactor
impl<Src, Dst, T: Clone + Div<T, Output=T>> Div<ScaleFactor<T, Src, Dst>> for Length<T, Dst> {
    type Output = Length<T, Src>;
    #[inline]
    fn div(self, scale: ScaleFactor<T, Src, Dst>) -> Length<T, Src> {
        Length::new(self.get() / scale.get())
    }
}

// -length
impl <U, T:Clone + Neg<Output=T>> Neg for Length<T, U> {
    type Output = Length<T, U>;
    #[inline]
    fn neg(self) -> Length<T, U> {
        Length::new(-self.get())
    }
}

impl<Unit, T0: NumCast + Clone> Length<T0, Unit> {
    /// Cast from one numeric representation to another, preserving the units.
    pub fn cast<T1: NumCast + Clone>(&self) -> Option<Length<T1, Unit>> {
        NumCast::from(self.get()).map(Length::new)
    }
}

impl<Unit, T: Clone + PartialEq> PartialEq for Length<T, Unit> {
    fn eq(&self, other: &Length<T, Unit>) -> bool { self.get().eq(&other.get()) }
}

impl<Unit, T: Clone + PartialOrd> PartialOrd for Length<T, Unit> {
    fn partial_cmp(&self, other: &Length<T, Unit>) -> Option<Ordering> {
        self.get().partial_cmp(&other.get())
    }
}

impl<Unit, T: Clone + Eq> Eq for Length<T, Unit> {}

impl<Unit, T: Clone + Ord> Ord for Length<T, Unit> {
    fn cmp(&self, other: &Length<T, Unit>) -> Ordering { self.get().cmp(&other.get()) }
}

impl<Unit, T: Zero> Zero for Length<T, Unit> {
    fn zero() -> Length<T, Unit> {
        Length::new(Zero::zero())
    }
}

#[cfg(test)]
mod tests {
    use super::Length;
    use num_traits::Saturating;
    use scale_factor::ScaleFactor;

    enum Inch {}
    enum Mm {}

    #[test]
    fn test_length() {
        let mm_per_inch: ScaleFactor<f32, Inch, Mm> = ScaleFactor::new(25.4);

        let one_foot: Length<f32, Inch> = Length::new(12.0);
        let two_feet = one_foot.clone() + one_foot.clone();
        let zero_feet = one_foot.clone() - one_foot.clone();

        assert_eq!(one_foot.get(), 12.0);
        assert_eq!(two_feet.get(), 24.0);
        assert_eq!(zero_feet.get(), 0.0);

        assert!(one_foot == one_foot);
        assert!(two_feet != one_foot);

        assert!(zero_feet <  one_foot);
        assert!(zero_feet <= one_foot);
        assert!(two_feet  >  one_foot);
        assert!(two_feet  >= one_foot);

        assert!(  two_feet <= two_feet);
        assert!(  two_feet >= two_feet);
        assert!(!(two_feet >  two_feet));
        assert!(!(two_feet <  two_feet));

        let one_foot_in_mm: Length<f32, Mm> = one_foot * mm_per_inch;

        assert_eq!(one_foot_in_mm, Length::new(304.8));
        assert_eq!(one_foot_in_mm / one_foot, mm_per_inch);

        let back_to_inches: Length<f32, Inch> = one_foot_in_mm / mm_per_inch;
        assert_eq!(one_foot, back_to_inches);

        let int_foot: Length<isize, Inch> = one_foot.cast().unwrap();
        assert_eq!(int_foot.get(), 12);

        let negative_one_foot = -one_foot;
        assert_eq!(negative_one_foot.get(), -12.0);

        let negative_two_feet = -two_feet;
        assert_eq!(negative_two_feet.get(), -24.0);

        let zero_feet: Length<f32, Inch> = Length::new(0.0);
        let negative_zero_feet = -zero_feet;
        assert_eq!(negative_zero_feet.get(), 0.0);
    }

    #[test]
    fn test_add() {
        let length1: Length<u8, Mm> = Length::new(250);
        let length2: Length<u8, Mm> = Length::new(5);

        let result = length1 + length2;

        assert_eq!(result.get(), 255);
    }

    #[test]
    fn test_saturating_add() {
        let length1: Length<u8, Mm> = Length::new(250);
        let length2: Length<u8, Mm> = Length::new(6);

        let result = length1.saturating_add(length2);

        assert_eq!(result.get(), 255);
    }

    #[test]
    fn test_saturating_sub() {
        let length1: Length<u8, Mm> = Length::new(5);
        let length2: Length<u8, Mm> = Length::new(10);

        let result = length1.saturating_sub(length2);

        assert_eq!(result.get(), 0);
    }

    #[test]
    fn test_addassign() {
        let one_cm: Length<f32, Mm> = Length::new(10.0);
        let mut measurement: Length<f32, Mm> = Length::new(5.0);

        measurement += one_cm;

        assert_eq!(measurement.get(), 15.0);
    }

    #[test]
    fn test_subassign() {
        let one_cm: Length<f32, Mm> = Length::new(10.0);
        let mut measurement: Length<f32, Mm> = Length::new(5.0);

        measurement -= one_cm;

        assert_eq!(measurement.get(), -5.0);
    }
}
<|endoftext|>


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

use anyhow::Result;
use mundane::hash::{Digest, Hasher, Sha256};
use std::time::Duration;
use thiserror::Error;
use zerocopy::AsBytes;

use crate::bootstrap::{BootstrapOptions, UpdateProtocol};
use crate::io::spi::{Target, Transfer};

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 2048 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const FLASH_BASE_ADDRESS: usize = 65536 * 8;

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::hash(&frame[Frame::HASH_LEN..]);
        sha.bytes()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let sha = Sha256::hash(self.as_bytes());
        let mut digest = sha.bytes();
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Vec<Frame> {
        let mut frames = Vec::new();

        let max_addr = (payload
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frame_num = 0;
        let mut addr = 0;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: (addr + Self::FLASH_BASE_ADDRESS) as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frame.header.hash = frame.header_hash();
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames
    }
}

/// Implements the bootstrap protocol of previous Google Titan family chips.
pub struct Legacy {
    pub inter_frame_delay: Duration,
    pub flash_erase_delay: Duration,
}

impl Legacy {
    const INTER_FRAME_DELAY: Duration = Duration::from_millis(20);
    const FLASH_ERASE_DELAY: Duration = Duration::from_millis(200);

    /// Creates a new `Primitive` protocol updater from `options`.
    pub fn new(options: &BootstrapOptions) -> Self {
        Self {
            inter_frame_delay: options.inter_frame_delay.unwrap_or(Self::INTER_FRAME_DELAY),
            flash_erase_delay: options.flash_erase_delay.unwrap_or(Self::FLASH_ERASE_DELAY),
        }
    }
}

#[derive(Debug, Error)]
pub enum LegacyBootstrapError {
    #[error("Boot rom not ready")]
    NotReady,
    #[error("Unknown boot rom error: {0}")]
    Unknown(u8),
    #[error("Boot rom error: NOREQUEST")]
    NoRequest,
    #[error("Boot rom error: NOMAGIC")]
    NoMagic,
    #[error("Boot rom error: TOOBIG")]
    TooBig,
    #[error("Boot rom error: TOOHIGH")]
    TooHigh,
    #[error("Boot rom error: NOALIGN")]
    NoAlign,
    #[error("Boot rom error: NOROUND")]
    NoRound,
    #[error("Boot rom error: BADKEY")]
    BadKey,
    #[error("Boot rom error: BADSTART")]
    BadStart,
    #[error("Boot rom error: NOWIPE")]
    NoWipe,
    #[error("Boot rom error: NOWIPE0")]
    NoWipe0,
    #[error("Boot rom error: NOWIPE1")]
    NoWipe1,
    #[error("Boot rom error: NOTEMPTY")]
    NotEmpty,
    #[error("Boot rom error: NOWRITE")]
    NoWrite,
    #[error("Boot rom error: BADADR")]
    BadAdr,
    #[error("Boot rom error: OVERFLOW")]
    Overflow,
}

impl From<u8> for LegacyBootstrapError {
    fn from(value: u8) -> LegacyBootstrapError {
        match value {
            // All zeroes or all ones means that the bootloader is not yet ready to respond.
            0 | 255 => LegacyBootstrapError::NotReady,
            // Other values represent particular errors.
            1 => LegacyBootstrapError::NoRequest,
            2 => LegacyBootstrapError::NoMagic,
            3 => LegacyBootstrapError::TooBig,
            4 => LegacyBootstrapError::TooHigh,
            5 => LegacyBootstrapError::NoAlign,
            6 => LegacyBootstrapError::NoRound,
            7 => LegacyBootstrapError::BadKey,
            8 => LegacyBootstrapError::BadStart,
            10 => LegacyBootstrapError::NoWipe,
            11 => LegacyBootstrapError::NoWipe0,
            12 => LegacyBootstrapError::NoWipe1,
            13 => LegacyBootstrapError::NotEmpty,
            14 => LegacyBootstrapError::NoWrite,
            15 => LegacyBootstrapError::BadAdr,
            16 => LegacyBootstrapError::Overflow,
            n => LegacyBootstrapError::Unknown(n),
        }
    }
}

impl UpdateProtocol for Legacy {
    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(&self, spi: &dyn Target, payload: &[u8]) -> Result<()> {
        let frames = Frame::from_payload(payload);

        for (i, frame) in frames.iter().enumerate() {
            let want_hash = frames[i.saturating_sub(1)].frame_hash();
            // Repeatedly transmit the frame, until receiving expected response.
            loop {
                // TODO: Introduce a progress callback, to allow common code to show progress bar.
                eprint!("{}.", i);

                // Important? Comment in spiflash.cc indicates that it seems to corrupt without a
                // 1ms delay.
                std::thread::sleep(self.inter_frame_delay);

                // Write the frame and read back the hash of the previous frame.
                let mut response = [0u8; std::mem::size_of::<Frame>()];
                spi.run_transaction(&mut [Transfer::Both(frame.as_bytes(), &mut response)])?;

                if response.iter().all(|&x| x == response[0]) {
                    // A response consisteing of all identical bytes is a status code.
                    match LegacyBootstrapError::from(response[0]) {
                        LegacyBootstrapError::NotReady => continue, // Retry sending same frame.
                        error => return Err(error.into()),
                    }
                }

                // Compare response with checksum of last block sent.
                if response[..Frame::HASH_LEN] == want_hash {
                    // Successful response, move on to the next frame.
                    break;
                }
                log::error!(
                    "Frame hash mismatch device:{:x?} != frame:{:x?}.",
                    &response[..Frame::HASH_LEN],
                    want_hash
                );
                // Retry sending same frame.
            }
        }
        eprintln!("success");
        Ok(())
    }
}

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

    const SIMPLE_BIN: &[u8; 2048] = include_bytes!("simple.bin");

    #[test]
    fn test_small_binary() -> Result<()> {
        let frames = Frame::from_payload(SIMPLE_BIN);

        assert_eq!(frames[0].header.frame_num, 0);
        assert_eq!(frames[0].header.flash_offset, 0x80000);
        assert_eq!(
            hex::encode(frames[0].header.hash),
            "4e31bfd8b3be32358f2235c0f241f3970de575fc6aca0564aa6bf30adaf33910"
        );

        assert_eq!(frames[1].header.frame_num, 0x8000_0001);
        assert_eq!(frames[1].header.flash_offset, 0x807d8);
        assert_eq!(
            hex::encode(frames[1].header.hash),
            "ff584c07bbb0a039934a660bd49b7812af8ee847d1e675d9aba71c11fab3cfcb"
        );
        Ok(())
    }
}
<commit_msg>[opentitantool] Fixes to bootstrap of legacy Titan<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::Result;
use mundane::hash::{Digest, Hasher, Sha256};
use std::time::Duration;
use thiserror::Error;
use zerocopy::AsBytes;

use crate::bootstrap::{BootstrapOptions, UpdateProtocol};
use crate::io::spi::{Target, Transfer};

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 2048 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const FLASH_BASE_ADDRESS: usize = 65536 * 8;

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::hash(&frame[Frame::HASH_LEN..]);
        sha.bytes()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let sha = Sha256::hash(self.as_bytes());
        let mut digest = sha.bytes();
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Vec<Frame> {
        let mut frames = Vec::new();

        let max_addr = (payload
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frame_num = 0;
        let mut addr = 0;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: (addr + Self::FLASH_BASE_ADDRESS) as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames.iter_mut().for_each(|f| f.header.hash = f.header_hash());
        frames
    }
}

/// Implements the bootstrap protocol of previous Google Titan family chips.
pub struct Legacy {
    /// Controls the time that CS is deasserted between frames.  With Dauntless, 20ms has been
    /// observed to be "too long", in that after the final frame was transmitted, the bootloader
    /// could have already stopped listening to SPI and started booting the newly flashed code,
    /// before OpenTitan tool could read the ACK of the final transaction.  1ms is what the
    /// spiflash.cc tool from cr50-utils uses.
    pub inter_frame_delay: Duration,
}

impl Legacy {
    const INTER_FRAME_DELAY: Duration = Duration::from_millis(1);
    const MAX_CONSECUTIVE_ERRORS: u32 = 100;

    /// Creates a new `Primitive` protocol updater from `options`.
    pub fn new(options: &BootstrapOptions) -> Self {
        Self {
            inter_frame_delay: options.inter_frame_delay.unwrap_or(Self::INTER_FRAME_DELAY),
        }
    }
}

#[derive(Debug, Error)]
pub enum LegacyBootstrapError {
    #[error("Boot rom not ready")]
    NotReady,
    #[error("Unknown boot rom error: {0}")]
    Unknown(u8),
    #[error("Boot rom error: NOREQUEST")]
    NoRequest,
    #[error("Boot rom error: NOMAGIC")]
    NoMagic,
    #[error("Boot rom error: TOOBIG")]
    TooBig,
    #[error("Boot rom error: TOOHIGH")]
    TooHigh,
    #[error("Boot rom error: NOALIGN")]
    NoAlign,
    #[error("Boot rom error: NOROUND")]
    NoRound,
    #[error("Boot rom error: BADKEY")]
    BadKey,
    #[error("Boot rom error: BADSTART")]
    BadStart,
    #[error("Boot rom error: NOWIPE")]
    NoWipe,
    #[error("Boot rom error: NOWIPE0")]
    NoWipe0,
    #[error("Boot rom error: NOWIPE1")]
    NoWipe1,
    #[error("Boot rom error: NOTEMPTY")]
    NotEmpty,
    #[error("Boot rom error: NOWRITE")]
    NoWrite,
    #[error("Boot rom error: BADADR")]
    BadAdr,
    #[error("Boot rom error: OVERFLOW")]
    Overflow,
    #[error("Repeated errors communicating with boot rom")]
    RepeatedErrors,
}

impl From<u8> for LegacyBootstrapError {
    fn from(value: u8) -> LegacyBootstrapError {
        match value {
            // All zeroes or all ones means that the bootloader is not yet ready to respond.
            0 | 255 => LegacyBootstrapError::NotReady,
            // Other values represent particular errors.
            1 => LegacyBootstrapError::NoRequest,
            2 => LegacyBootstrapError::NoMagic,
            3 => LegacyBootstrapError::TooBig,
            4 => LegacyBootstrapError::TooHigh,
            5 => LegacyBootstrapError::NoAlign,
            6 => LegacyBootstrapError::NoRound,
            7 => LegacyBootstrapError::BadKey,
            8 => LegacyBootstrapError::BadStart,
            10 => LegacyBootstrapError::NoWipe,
            11 => LegacyBootstrapError::NoWipe0,
            12 => LegacyBootstrapError::NoWipe1,
            13 => LegacyBootstrapError::NotEmpty,
            14 => LegacyBootstrapError::NoWrite,
            15 => LegacyBootstrapError::BadAdr,
            16 => LegacyBootstrapError::Overflow,
            n => LegacyBootstrapError::Unknown(n),
        }
    }
}

impl UpdateProtocol for Legacy {
    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(&self, spi: &dyn Target, payload: &[u8]) -> Result<()> {
        let frames = Frame::from_payload(payload);

        // All frames up to but not including this index have been ack'ed by the bootloader.
        // Once this reaches frames.len(), the operation was successful.
        let mut first_unacked_index = 0;

        // Counts the number of repeated failures in getting a frame ack.
        let mut consecutive_errors = 0;

        // Set if the past transaction ack'ed the block sent before that one, while transmitting
        // the next.  If so, we heuristically assume that the block just sent will be ack'ed in
        // the upcoming bidirectional transaction.
        let mut optimistic = false;

        // Set if the past transaction ack'ed the block sent before that one, while
        // re-transmitting the same block.  If so, we do not care about the ack in the upcoming
        // bidirectional transaction, as it turns out it was not necessary to perform the most
        // recent re-transmission.
        let mut becoming_optimistic = false;

        loop {
            if consecutive_errors > Self::MAX_CONSECUTIVE_ERRORS {
                return Err(LegacyBootstrapError::RepeatedErrors.into())
            }

            let second_unacked_index = (first_unacked_index + 1).min(frames.len() - 1);
            let transmit_index =
                if optimistic { second_unacked_index } else { first_unacked_index };
            let frame = &frames[transmit_index];
            eprint!("{}.", transmit_index);
            std::thread::sleep(self.inter_frame_delay);

            // Write the frame and read back the ack of a previously transmitted frame.
            let mut response = [0u8; std::mem::size_of::<Frame>()];
            spi.run_transaction(&mut [Transfer::Both(frame.as_bytes(), &mut response)])?;

            if becoming_optimistic {
                optimistic = true;
                becoming_optimistic = false;
                continue;
            }

            if response.iter().all(|&x| x == response[0]) {
                // A response consisteing of all identical bytes is a status code.
                match LegacyBootstrapError::from(response[0]) {
                    LegacyBootstrapError::NotReady => {
                        consecutive_errors += 1;
                        continue; // Retry sending same frame.
                    }
                    error => return Err(error.into()),
                }
            }

            if response[..Frame::HASH_LEN] == frames[first_unacked_index].frame_hash() {
                first_unacked_index = first_unacked_index + 1;
            } else if response[..Frame::HASH_LEN] == frames[second_unacked_index].frame_hash() {
                first_unacked_index = second_unacked_index + 1;
            } else {
                consecutive_errors += 1;
                optimistic = false;
                continue;
            }

            consecutive_errors = 0;
            if first_unacked_index == frames.len() {
                // All frames acked, we are done.
                break;
            }
            if !optimistic {
                // Heuristic, become optimistic after second or later frame is ack'ed.
                becoming_optimistic = first_unacked_index > 1;
            }
        }
        eprintln!("success");
        Ok(())
    }
}

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

    const SIMPLE_BIN: &[u8; 2048] = include_bytes!("simple.bin");

    #[test]
    fn test_small_binary() -> Result<()> {
        let frames = Frame::from_payload(SIMPLE_BIN);

        assert_eq!(frames[0].header.frame_num, 0);
        assert_eq!(frames[0].header.flash_offset, 0x80000);
        assert_eq!(
            hex::encode(frames[0].header.hash),
            "4e31bfd8b3be32358f2235c0f241f3970de575fc6aca0564aa6bf30adaf33910"
        );

        assert_eq!(frames[1].header.frame_num, 0x8000_0001);
        assert_eq!(frames[1].header.flash_offset, 0x807d8);
        assert_eq!(
            hex::encode(frames[1].header.hash),
            "ff584c07bbb0a039934a660bd49b7812af8ee847d1e675d9aba71c11fab3cfcb"
        );
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::error::Error;
use std::fmt;
use std::io::{self, Cursor};
use std::str;

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum BEncode {
    Int(i64),
    String(Vec<u8>),
    List(Vec<BEncode>),
    Dict(BTreeMap<String, BEncode>),
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum BError {
    UTF8Decode,
    InvalidDict,
    InvalidChar(u8),
    ParseInt,
    EOF,
    IO,
}

impl fmt::Display for BError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            BError::UTF8Decode => write!(f, "UTF8 Decoding Error"),
            BError::InvalidDict => write!(f, "Invalid BEncoded dictionary"),
            BError::InvalidChar(c) => write!(f, "Invalid character: {}", char::from(c)),
            BError::ParseInt => write!(f, "Invalid integer value encountered"),
            BError::EOF => write!(f, "Unexpected EOF in data"),
            BError::IO => write!(f, "IO error"),
        }
    }
}

impl Error for BError {
    fn description(&self) -> &str {
        "BEncode processing error"
    }
}

impl BEncode {
    pub fn from_int(i: i64) -> BEncode {
        BEncode::Int(i)
    }

    #[cfg_attr(feature = "cargo-clippy", allow(should_implement_trait))]
    pub fn from_str(s: &str) -> BEncode {
        BEncode::String(Vec::from(s))
    }

    pub fn into_int(self) -> Option<i64> {
        match self {
            BEncode::Int(v) => Some(v),
            _ => None,
        }
    }

    pub fn into_bytes(self) -> Option<Vec<u8>> {
        match self {
            BEncode::String(v) => Some(v),
            _ => None,
        }
    }

    pub fn into_string(self) -> Option<String> {
        match self {
            BEncode::String(v) => String::from_utf8(v).ok(),
            _ => None,
        }
    }

    pub fn into_list(self) -> Option<Vec<BEncode>> {
        match self {
            BEncode::List(v) => Some(v),
            _ => None,
        }
    }

    pub fn into_dict(self) -> Option<BTreeMap<String, BEncode>> {
        match self {
            BEncode::Dict(v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_int(&self) -> Option<&i64> {
        match *self {
            BEncode::Int(ref v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_bytes(&self) -> Option<&Vec<u8>> {
        match *self {
            BEncode::String(ref v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_str(&self) -> Option<&str> {
        match *self {
            BEncode::String(ref v) => str::from_utf8(v).ok(),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_list(&self) -> Option<&Vec<BEncode>> {
        match *self {
            BEncode::List(ref v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_dict(&self) -> Option<&BTreeMap<String, BEncode>> {
        match *self {
            BEncode::Dict(ref v) => Some(v),
            _ => None,
        }
    }

    pub fn encode_to_buf(&self) -> Vec<u8> {
        let mut buf = Cursor::new(Vec::new());
        self.encode(&mut buf).unwrap();
        buf.into_inner()
    }

    pub fn encode<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
        enum Token<'a> {
            B(&'a BEncode),
            OS(&'a str),
            E,
        }

        let mut toks = vec![Token::B(self)];
        while let Some(tok) = toks.pop() {
            match tok {
                Token::B(&BEncode::Int(i)) => {
                    write!(w, "i{}e", i)?;
                }
                Token::B(&BEncode::String(ref s)) => {
                    write!(w, "{}:", s.len())?;
                    w.write_all(s)?;
                }
                Token::B(&BEncode::List(ref v)) => {
                    write!(w, "l")?;
                    toks.push(Token::E);
                    toks.extend(v.iter().rev().map(|v| Token::B(v)));
                }
                Token::B(&BEncode::Dict(ref d)) => {
                    write!(w, "d")?;
                    toks.push(Token::E);
                    for (k, v) in d.iter().rev() {
                        toks.push(Token::B(v));
                        toks.push(Token::OS(k));
                    }
                }
                Token::OS(s) => {
                    write!(w, "{}:", s.len())?;
                    w.write_all(s.as_bytes())?;
                }
                Token::E => {
                    write!(w, "e")?;
                }
            }
        }
        Ok(())
    }
}

pub fn decode_buf(bytes: &[u8]) -> Result<BEncode, BError> {
    decode(&mut Cursor::new(bytes))
}

pub fn decode_buf_first(bytes: &[u8]) -> Result<BEncode, BError> {
    decode_first(&mut Cursor::new(bytes))
}

pub fn decode_first<R: io::Read>(bytes: &mut R) -> Result<BEncode, BError> {
    do_decode(bytes, true)
}

pub fn decode<R: io::Read>(bytes: &mut R) -> Result<BEncode, BError> {
    do_decode(bytes, false)
}

fn do_decode<R: io::Read>(bytes: &mut R, first: bool) -> Result<BEncode, BError> {
    enum Kind {
        Dict(usize),
        List(usize),
    }
    let mut cstack = vec![];
    let mut vstack = vec![];
    let mut buf = [0];
    while !first || !(cstack.is_empty() && vstack.len() == 1) {
        match next_byte(bytes, &mut buf) {
            Ok(b'i') => {
                // Multiple non complex values are not allowed
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                let s = read_until(bytes, b'e', &mut buf)?;
                vstack.push(BEncode::Int(decode_int(s)?));
            }
            Ok(b'l') => {
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                cstack.push(Kind::List(vstack.len()));
            }
            Ok(b'd') => {
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                cstack.push(Kind::Dict(vstack.len()));
            }
            Err(BError::EOF) => break,
            Ok(b'e') => match cstack.pop() {
                Some(Kind::List(i)) => {
                    let mut l = Vec::with_capacity(vstack.len() - i);
                    while vstack.len() > i {
                        l.push(vstack.pop().unwrap());
                    }
                    l.reverse();
                    vstack.push(BEncode::List(l));
                }
                Some(Kind::Dict(i)) => {
                    let mut d = BTreeMap::new();
                    if (vstack.len() - i) % 2 != 0 {
                        return Err(BError::InvalidDict);
                    }
                    while vstack.len() > i {
                        let val = vstack.pop().unwrap();
                        match vstack.pop().and_then(BEncode::into_string) {
                            Some(key) => {
                                d.insert(key, val);
                            }
                            None => return Err(BError::InvalidDict),
                        }
                    }
                    vstack.push(BEncode::Dict(d))
                }
                None => return Err(BError::InvalidChar(b'e')),
            },
            Ok(d @ b'0'..=b'9') => {
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                let mut slen = read_until(bytes, b':', &mut buf)?;
                slen.insert(0, d);
                let len = decode_int(slen)?;
                let mut v = vec![0u8; len as usize];
                bytes.read_exact(&mut v).map_err(|_| BError::EOF)?;
                vstack.push(BEncode::String(v));
            }
            Err(e) => return Err(e),
            Ok(c) => return Err(BError::InvalidChar(c)),
        }
    }

    if cstack.is_empty() && vstack.len() == 1 {
        Ok(vstack.into_iter().next().unwrap())
    } else {
        Err(BError::EOF)
    }
}

fn next_byte<R: io::Read>(r: &mut R, buf: &mut [u8; 1]) -> Result<u8, BError> {
    let amnt = r.read(buf).map_err(|_| BError::IO)?;
    if amnt == 0 {
        Err(BError::EOF)
    } else {
        Ok(buf[0])
    }
}

fn read_until<R: io::Read>(r: &mut R, b: u8, buf: &mut [u8; 1]) -> Result<Vec<u8>, BError> {
    let mut v = vec![];
    loop {
        let n = next_byte(r, buf)?;
        if b == n {
            return Ok(v);
        }
        v.push(n);
    }
}

fn decode_int(v: Vec<u8>) -> Result<i64, BError> {
    String::from_utf8(v)
        .map_err(|_| BError::UTF8Decode)
        .and_then(|i| i.parse().map_err(|_| BError::ParseInt))
}

#[cfg(test)]
mod tests {
    use super::{decode_buf, decode_buf_first, BEncode};
    use std::collections::BTreeMap;

    #[test]
    fn test_encode_decode() {
        let i = BEncode::Int(-10);
        let mut v = Vec::new();
        i.encode(&mut v).unwrap();
        assert_eq!(v, b"i-10e");

        let s = BEncode::String(Vec::from(&b"asdf"[..]));
        v = Vec::new();
        s.encode(&mut v).unwrap();
        assert_eq!(v, b"4:asdf");

        let s2r = [1u8, 2, 3, 4];
        let s2e = [52u8, 58, 1, 2, 3, 4];
        let s2 = BEncode::String(Vec::from(&s2r[..]));
        v = Vec::new();
        s2.encode(&mut v).unwrap();
        assert_eq!(v, &s2e);

        let l = BEncode::List(vec![i.clone(), s.clone()]);
        v = Vec::new();
        l.encode(&mut v).unwrap();
        assert_eq!(v, b"li-10e4:asdfe");

        let mut map = BTreeMap::new();
        map.insert(String::from("asdf"), i.clone());
        map.insert(String::from("qwerty"), i.clone());
        let d = BEncode::Dict(map);
        v = Vec::new();
        d.encode(&mut v).unwrap();
        println!("{:?}", std::str::from_utf8(&v));
        assert_eq!(v, b"d4:asdfi-10e6:qwertyi-10ee");

        decode_encode(b"d4:asdfi-10e6:qwertyi-10ee");

        decode_encode(b"d1:rd2:id20:mnopqrstuvwxyz123456e1:t2:aa1:y1:re");

        encode_decode(&i);
        encode_decode(&s);
        encode_decode(&s2);
        encode_decode(&l);
        encode_decode(&d);
    }

    #[test]
    fn test_first_valid() {
        let twoint = b"i123ei123e";
        assert!(decode_buf_first(twoint).is_ok());
    }

    #[test]
    fn test_invalid() {
        let badint = b"i-123.4e";
        let twoint = b"i123ei123e";
        let badstr = b"5:eeeeeeeeeeee";
        let badstr2 = b"5:e";
        let badstr3 = b"-1:e";
        let badstr4 = b"1:a2:ab";
        let badlist = b"l123e";
        let badlist2 = b"li123e";
        let badlist3 = b"lllllllllllllllllllllllllllllllllllllleeeeeeeeeeee";
        let badlist4 = b"lele";
        let baddict = b"d1:ae";
        let baddict2 = b"di123ei123ee";
        assert!(decode_buf(badint).is_err());
        assert!(decode_buf(twoint).is_err());
        assert!(decode_buf(badstr).is_err());
        assert!(decode_buf(badstr2).is_err());
        assert!(decode_buf(badstr3).is_err());
        assert!(decode_buf(badstr4).is_err());
        assert!(decode_buf(badlist).is_err());
        assert!(decode_buf(badlist2).is_err());
        assert!(decode_buf(badlist3).is_err());
        assert!(decode_buf(badlist4).is_err());
        assert!(decode_buf(baddict).is_err());
        assert!(decode_buf(baddict2).is_err());
    }

    fn encode_decode(b: &BEncode) {
        let mut v = Vec::new();
        b.encode(&mut v).unwrap();
        assert_eq!(b, &decode_buf(&v).unwrap());
    }

    fn decode_encode(d: &[u8]) {
        let mut v = Vec::new();
        let dec = decode_buf(d).unwrap();
        dec.encode(&mut v).unwrap();
        assert_eq!(d, &v[..]);
    }
}
<commit_msg>Prevent OOM from int parsing in bencode library<commit_after>use std::collections::BTreeMap;
use std::error::Error;
use std::{cmp, fmt, str};
use std::io::{self, Cursor};

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum BEncode {
    Int(i64),
    String(Vec<u8>),
    List(Vec<BEncode>),
    Dict(BTreeMap<String, BEncode>),
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum BError {
    UTF8Decode,
    InvalidDict,
    InvalidChar(u8),
    ParseInt,
    EOF,
    IO,
}

/// This controls the maximum allocation size we'll perform
/// at once. Needed for parsing strings without OOMing
const MAX_ALLOC_LEN: usize = 4 * 1024 * 1024;

impl fmt::Display for BError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            BError::UTF8Decode => write!(f, "UTF8 Decoding Error"),
            BError::InvalidDict => write!(f, "Invalid BEncoded dictionary"),
            BError::InvalidChar(c) => write!(f, "Invalid character: {}", char::from(c)),
            BError::ParseInt => write!(f, "Invalid integer value encountered"),
            BError::EOF => write!(f, "Unexpected EOF in data"),
            BError::IO => write!(f, "IO error"),
        }
    }
}

impl Error for BError {
    fn description(&self) -> &str {
        "BEncode processing error"
    }
}

impl BEncode {
    pub fn from_int(i: i64) -> BEncode {
        BEncode::Int(i)
    }

    #[cfg_attr(feature = "cargo-clippy", allow(should_implement_trait))]
    pub fn from_str(s: &str) -> BEncode {
        BEncode::String(Vec::from(s))
    }

    pub fn into_int(self) -> Option<i64> {
        match self {
            BEncode::Int(v) => Some(v),
            _ => None,
        }
    }

    pub fn into_bytes(self) -> Option<Vec<u8>> {
        match self {
            BEncode::String(v) => Some(v),
            _ => None,
        }
    }

    pub fn into_string(self) -> Option<String> {
        match self {
            BEncode::String(v) => String::from_utf8(v).ok(),
            _ => None,
        }
    }

    pub fn into_list(self) -> Option<Vec<BEncode>> {
        match self {
            BEncode::List(v) => Some(v),
            _ => None,
        }
    }

    pub fn into_dict(self) -> Option<BTreeMap<String, BEncode>> {
        match self {
            BEncode::Dict(v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_int(&self) -> Option<&i64> {
        match *self {
            BEncode::Int(ref v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_bytes(&self) -> Option<&Vec<u8>> {
        match *self {
            BEncode::String(ref v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_str(&self) -> Option<&str> {
        match *self {
            BEncode::String(ref v) => str::from_utf8(v).ok(),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_list(&self) -> Option<&Vec<BEncode>> {
        match *self {
            BEncode::List(ref v) => Some(v),
            _ => None,
        }
    }

    #[allow(dead_code)]
    pub fn as_dict(&self) -> Option<&BTreeMap<String, BEncode>> {
        match *self {
            BEncode::Dict(ref v) => Some(v),
            _ => None,
        }
    }

    pub fn encode_to_buf(&self) -> Vec<u8> {
        let mut buf = Cursor::new(Vec::new());
        self.encode(&mut buf).unwrap();
        buf.into_inner()
    }

    pub fn encode<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
        enum Token<'a> {
            B(&'a BEncode),
            OS(&'a str),
            E,
        }

        let mut toks = vec![Token::B(self)];
        while let Some(tok) = toks.pop() {
            match tok {
                Token::B(&BEncode::Int(i)) => {
                    write!(w, "i{}e", i)?;
                }
                Token::B(&BEncode::String(ref s)) => {
                    write!(w, "{}:", s.len())?;
                    w.write_all(s)?;
                }
                Token::B(&BEncode::List(ref v)) => {
                    write!(w, "l")?;
                    toks.push(Token::E);
                    toks.extend(v.iter().rev().map(|v| Token::B(v)));
                }
                Token::B(&BEncode::Dict(ref d)) => {
                    write!(w, "d")?;
                    toks.push(Token::E);
                    for (k, v) in d.iter().rev() {
                        toks.push(Token::B(v));
                        toks.push(Token::OS(k));
                    }
                }
                Token::OS(s) => {
                    write!(w, "{}:", s.len())?;
                    w.write_all(s.as_bytes())?;
                }
                Token::E => {
                    write!(w, "e")?;
                }
            }
        }
        Ok(())
    }
}

pub fn decode_buf(bytes: &[u8]) -> Result<BEncode, BError> {
    decode(&mut Cursor::new(bytes))
}

pub fn decode_buf_first(bytes: &[u8]) -> Result<BEncode, BError> {
    decode_first(&mut Cursor::new(bytes))
}

pub fn decode_first<R: io::Read>(bytes: &mut R) -> Result<BEncode, BError> {
    do_decode(bytes, true)
}

pub fn decode<R: io::Read>(bytes: &mut R) -> Result<BEncode, BError> {
    do_decode(bytes, false)
}

fn do_decode<R: io::Read>(bytes: &mut R, first: bool) -> Result<BEncode, BError> {
    enum Kind {
        Dict(usize),
        List(usize),
    }
    let mut cstack = vec![];
    let mut vstack = vec![];
    let mut buf = [0];
    while !first || !(cstack.is_empty() && vstack.len() == 1) {
        match next_byte(bytes, &mut buf) {
            Ok(b'i') => {
                // Multiple non complex values are not allowed
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                let s = read_until(bytes, b'e', &mut buf)?;
                vstack.push(BEncode::Int(decode_int(s)?));
            }
            Ok(b'l') => {
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                cstack.push(Kind::List(vstack.len()));
            }
            Ok(b'd') => {
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                cstack.push(Kind::Dict(vstack.len()));
            }
            Err(BError::EOF) => break,
            Ok(b'e') => match cstack.pop() {
                Some(Kind::List(i)) => {
                    let mut l = Vec::with_capacity(vstack.len() - i);
                    while vstack.len() > i {
                        l.push(vstack.pop().unwrap());
                    }
                    l.reverse();
                    vstack.push(BEncode::List(l));
                }
                Some(Kind::Dict(i)) => {
                    let mut d = BTreeMap::new();
                    if (vstack.len() - i) % 2 != 0 {
                        return Err(BError::InvalidDict);
                    }
                    while vstack.len() > i {
                        let val = vstack.pop().unwrap();
                        match vstack.pop().and_then(BEncode::into_string) {
                            Some(key) => {
                                d.insert(key, val);
                            }
                            None => return Err(BError::InvalidDict),
                        }
                    }
                    vstack.push(BEncode::Dict(d))
                }
                None => return Err(BError::InvalidChar(b'e')),
            },
            Ok(d @ b'0'..=b'9') => {
                if cstack.is_empty() && !vstack.is_empty() {
                    return Err(BError::EOF);
                }
                let mut slen = read_until(bytes, b':', &mut buf)?;
                slen.insert(0, d);
                let len = decode_int(slen)?;
                let mut v = vec![];
                while v.len() < len as usize {
                    let to_read = cmp::min(MAX_ALLOC_LEN, len as usize - v.len());
                    v.resize(v.len() + to_read, 0u8);
                    let read_start = v.len() - to_read;
                    let read_end = v.len();
                    bytes.read_exact(&mut v[read_start..read_end])
                        .map_err(|_| BError::EOF)?;
                }
                vstack.push(BEncode::String(v));
            }
            Err(e) => return Err(e),
            Ok(c) => return Err(BError::InvalidChar(c)),
        }
    }

    if cstack.is_empty() && vstack.len() == 1 {
        Ok(vstack.into_iter().next().unwrap())
    } else {
        Err(BError::EOF)
    }
}

fn next_byte<R: io::Read>(r: &mut R, buf: &mut [u8; 1]) -> Result<u8, BError> {
    let amnt = r.read(buf).map_err(|_| BError::IO)?;
    if amnt == 0 {
        Err(BError::EOF)
    } else {
        Ok(buf[0])
    }
}

fn read_until<R: io::Read>(r: &mut R, b: u8, buf: &mut [u8; 1]) -> Result<Vec<u8>, BError> {
    let mut v = vec![];
    loop {
        let n = next_byte(r, buf)?;
        if b == n {
            return Ok(v);
        }
        v.push(n);
    }
}

fn decode_int(v: Vec<u8>) -> Result<i64, BError> {
    String::from_utf8(v)
        .map_err(|_| BError::UTF8Decode)
        .and_then(|i| i.parse().map_err(|_| BError::ParseInt))
}

#[cfg(test)]
mod tests {
    use super::{decode_buf, decode_buf_first, BEncode};
    use std::collections::BTreeMap;

    #[test]
    fn test_encode_decode() {
        let i = BEncode::Int(-10);
        let mut v = Vec::new();
        i.encode(&mut v).unwrap();
        assert_eq!(v, b"i-10e");

        let s = BEncode::String(Vec::from(&b"asdf"[..]));
        v = Vec::new();
        s.encode(&mut v).unwrap();
        assert_eq!(v, b"4:asdf");

        let s2r = [1u8, 2, 3, 4];
        let s2e = [52u8, 58, 1, 2, 3, 4];
        let s2 = BEncode::String(Vec::from(&s2r[..]));
        v = Vec::new();
        s2.encode(&mut v).unwrap();
        assert_eq!(v, &s2e);

        let l = BEncode::List(vec![i.clone(), s.clone()]);
        v = Vec::new();
        l.encode(&mut v).unwrap();
        assert_eq!(v, b"li-10e4:asdfe");

        let mut map = BTreeMap::new();
        map.insert(String::from("asdf"), i.clone());
        map.insert(String::from("qwerty"), i.clone());
        let d = BEncode::Dict(map);
        v = Vec::new();
        d.encode(&mut v).unwrap();
        println!("{:?}", std::str::from_utf8(&v));
        assert_eq!(v, b"d4:asdfi-10e6:qwertyi-10ee");

        decode_encode(b"d4:asdfi-10e6:qwertyi-10ee");

        decode_encode(b"d1:rd2:id20:mnopqrstuvwxyz123456e1:t2:aa1:y1:re");

        encode_decode(&i);
        encode_decode(&s);
        encode_decode(&s2);
        encode_decode(&l);
        encode_decode(&d);
    }

    #[test]
    fn test_first_valid() {
        let twoint = b"i123ei123e";
        assert!(decode_buf_first(twoint).is_ok());
    }

    #[test]
    fn test_invalid() {
        let badint = b"i-123.4e";
        let badint_oom = b"7777777777777:";
        let twoint = b"i123ei123e";
        let badstr = b"5:eeeeeeeeeeee";
        let badstr2 = b"5:e";
        let badstr3 = b"-1:e";
        let badstr4 = b"1:a2:ab";
        let badlist = b"l123e";
        let badlist2 = b"li123e";
        let badlist3 = b"lllllllllllllllllllllllllllllllllllllleeeeeeeeeeee";
        let badlist4 = b"lele";
        let baddict = b"d1:ae";
        let baddict2 = b"di123ei123ee";
        assert!(decode_buf(badint).is_err());
        assert!(decode_buf(badint_oom).is_err());
        assert!(decode_buf(twoint).is_err());
        assert!(decode_buf(badstr).is_err());
        assert!(decode_buf(badstr2).is_err());
        assert!(decode_buf(badstr3).is_err());
        assert!(decode_buf(badstr4).is_err());
        assert!(decode_buf(badlist).is_err());
        assert!(decode_buf(badlist2).is_err());
        assert!(decode_buf(badlist3).is_err());
        assert!(decode_buf(badlist4).is_err());
        assert!(decode_buf(baddict).is_err());
        assert!(decode_buf(baddict2).is_err());
    }

    fn encode_decode(b: &BEncode) {
        let mut v = Vec::new();
        b.encode(&mut v).unwrap();
        assert_eq!(b, &decode_buf(&v).unwrap());
    }

    fn decode_encode(d: &[u8]) {
        let mut v = Vec::new();
        let dec = decode_buf(d).unwrap();
        dec.encode(&mut v).unwrap();
        assert_eq!(d, &v[..]);
    }
}
<|endoftext|>