<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//!Media module
//!
//!Provides functionality to upload images, GIFs and videos using Twitter Media API

use std::borrow::Cow;
use std::collections::HashMap;

use futures::{Future, Async, Poll};
use rustc_serialize::{json, base64};
use rustc_serialize::base64::{ToBase64};

use common::*;
use error;
use error::Error::{InvalidResponse, MissingValue};
use links;
use auth;

use mime;

/// A collection of convenience functions that return media types accepted by Twitter.
pub mod media_types {
    use mime::{self, Mime};

    /// PNG images.
    pub fn image_png() -> Mime {
        mime::IMAGE_PNG
    }

    /// JPG images.
    pub fn image_jpg() -> Mime {
        mime::IMAGE_JPEG
    }

    /// WEBP images.
    pub fn image_webp() -> Mime {
        "image/webp".parse().unwrap()
    }

    /// GIF images, both animated and static.
    pub fn image_gif() -> Mime {
        mime::IMAGE_GIF
    }

    /// MP4 videos.
    pub fn video_mp4() -> Mime {
        "video/mp4".parse().unwrap()
    }
}

#[derive(Debug)]
///Media's upload progressing info.
pub enum ProgressInfo {
    ///Video is pending for processing. Contains number of seconds after which to check.
    Pending(u64),
    ///Video is beeing processed. Contains number of seconds after which to check.
    InProgress(u64),
    ///Video's processing failed. Contains reason.
    Failed(String),
    ///Video's processing is finished. Media can be used in other API calls.
    Success
}

impl FromJson for ProgressInfo {
    fn from_json(input: &json::Json) -> Result<Self, error::Error> {
        field_present!(input, state);
        let state: String = try!(field(input, "state"));

        match state.as_ref() {
            "pending" => {
                field_present!(input, check_after_secs);
                Ok(ProgressInfo::Pending(try!(field(input, "check_after_secs"))))
            },
            "in_progress" => {
                field_present!(input, check_after_secs);
                Ok(ProgressInfo::InProgress(try!(field(input, "check_after_secs"))))
            },
            "failed" => {
                let error = try!(input.find("error").ok_or(MissingValue("error")));

                let name = try!(error.find("name").ok_or(MissingValue("name")));
                let message = try!(error.find("message").ok_or(MissingValue("message")));

                let name = try!(name.as_string().ok_or(InvalidResponse("Expected string in error's name", None)));
                let message = try!(message.as_string().ok_or(InvalidResponse("Expected string in error's message", None)));
                Ok(ProgressInfo::Failed(format!("{}: {}", name, message)))
            },
            "succeeded" => Ok(ProgressInfo::Success),
            state => Err(InvalidResponse("Unexpected progress info state", Some(state.to_string())))

        }
    }
}

///Represents media file that is uploaded on twitter.
pub struct Media {
    ///ID that can be used in API calls (e.g. attach to tweet).
    pub id: u64,
    ///Number of second the media can be used in other API calls.
    pub expires_after: u64,
    ///Progress information. If present determines whether Media can be used.
    pub progress: Option<ProgressInfo>
}

impl FromJson for Media {
    fn from_json(input: &json::Json) -> Result<Self, error::Error> {
        if !input.is_object() {
            return Err(InvalidResponse("Tweet received json that wasn't an object", Some(input.to_string())));
        }

        field_present!(input, media_id);

        Ok(Media {
            id: try!(field(input, "media_id")),
            //We can miss this field on failed upload in which case 0 is pretty reasonable value.
            expires_after: field(input, "expires_after_secs").unwrap_or(0),
            progress: try!(field(input, "processing_info"))
        })
    }
}

///Uploads image using old twitter API.
///
///The image should be raw binary with content of file.
pub fn upload_image(image: &[u8], token: &auth::Token, handle: &Handle) -> FutureResponse<Media> {
    let mut params = HashMap::new();

    let config = base64::Config {
        char_set: base64::CharacterSet::Standard,
        newline: base64::Newline::LF,
        pad: true,
        line_length: None,
    };

    add_param(&mut params, "media_data", image.to_base64(config));

    let req = auth::post(links::medias::UPLOAD, token, Some(&params));
    make_parsed_future(handle, req)
}

/// Represents the kinda of media that Twitter will accept.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum MediaCategory {
    /// Static image. Four can be attached to a single tweet.
    Image,
    /// Animated GIF.
    Gif,
    /// Video.
    Video,
}

/// `Display` impl for `MediaCategory` so that `.to_string()` will return a string suitable for use
/// in an API call. This will turn the enum into `"tweet_image"`, `"tweet_gif"`, and
/// `"tweet_video"`.
impl ::std::fmt::Display for MediaCategory {
    fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match *self {
            MediaCategory::Image => write!(fmt, "tweet_image"),
            MediaCategory::Gif => write!(fmt, "tweet_gif"),
            MediaCategory::Video => write!(fmt, "tweet_video"),
        }
    }
}

/// A builder struct that allows you to build up parameters to a media upload before initiating it.
pub struct UploadBuilder<'a> {
    data: Cow<'a, [u8]>,
    media_type: mime::Mime,
    chunk_size: Option<usize>,
    category: Option<MediaCategory>,
}

impl<'a> UploadBuilder<'a> {
    /// Begins setting up a media upload call, with the given data and media type.
    ///
    /// For convenience functions to get known `media_type`s that Twitter will accept, see the
    /// [`media_types`] module.
    ///
    /// [`media_types`]: media_types/index.html
    pub fn new<V: Into<Cow<'a, [u8]>>>(data: V, media_type: mime::Mime) -> UploadBuilder<'a> {
        UploadBuilder {
            data: data.into(),
            media_type,
            chunk_size: None,
            category: None,
        }
    }

    /// Sets how many bytes to upload in one network call. By default this is set to 512 KiB.
    pub fn chunk_size(self, chunk_size: usize) -> Self {
        UploadBuilder {
            chunk_size: Some(chunk_size),
            ..self
        }
    }

    /// Sets the `media_category` sent to Twitter. When unset, it behaves as if you sent
    /// `MediaCategory::Image`.
    pub fn category(self, category: MediaCategory) -> Self {
        UploadBuilder {
            category: Some(category),
            ..self
        }
    }

    /// Collects the built-up parameters and begins the chunked upload.
    pub fn call(self, token: &auth::Token, handle: &Handle) -> UploadFuture<'a> {
        let mut params = HashMap::new();

        add_param(&mut params, "command", "INIT");
        add_param(&mut params, "total_bytes", self.data.len().to_string());
        add_param(&mut params, "media_type", self.media_type.to_string());

        if let Some(category) = self.category {
            add_param(&mut params, "media_category", category.to_string());
        }

        let req = auth::post(links::medias::UPLOAD, token, Some(&params));
        let loader = make_parsed_future(handle, req);
        UploadFuture {
            data: self.data,
            token: token.clone(),
            handle: handle.clone(),
            chunk_size: self.chunk_size.unwrap_or(1024 * 512), // 512 KiB default
            status: UploadInner::WaitingForInit(loader),
        }
    }
}

/// A `Future` that represents an in-progress media upload.
pub struct UploadFuture<'a> {
    data: Cow<'a, [u8]>,
    token: auth::Token,
    handle: Handle,
    chunk_size: usize,
    status: UploadInner,
}

/// The current status of an `UploadFuture`.
enum UploadInner {
    /// The `UploadFuture` is waiting to initialize the media upload session.
    WaitingForInit(FutureResponse<Media>),
    /// The `UploadFuture` is in the progress of uploading data.
    UploadingChunk(u64, usize, FutureResponse<()>),
    /// The `UploadFuture` is currently finalizing the media with Twitter.
    Finalizing(FutureResponse<Media>),
    /// The `UploadFuture` has completed, or has encountered an error.
    Invalid,
}

impl<'a> UploadFuture<'a> {
    fn get_chunk(&self, chunk_num: usize) -> Option<&[u8]> {
        let start = chunk_num * self.chunk_size;
        let end = (chunk_num + 1) * self.chunk_size;

        if start >= self.data.len() {
            None
        } else if end >= self.data.len() {
            Some(&self.data[start..])
        } else {
            Some(&self.data[start..end])
        }
    }

    fn append(&self, chunk_num: usize, media_id: u64) -> Option<FutureResponse<()>> {
        let mut chunk = self.get_chunk(chunk_num);
        if chunk.is_none() && chunk_num == 0 {
            chunk = Some(&[][..]);
        }

        if let Some(chunk) = chunk {
            let mut params = HashMap::new();

            let config = base64::Config {
                char_set: base64::CharacterSet::Standard,
                newline: base64::Newline::LF,
                pad: true,
                line_length: None,
            };

            add_param(&mut params, "command", "APPEND");
            add_param(&mut params, "media_id", media_id.to_string());
            add_param(&mut params, "media_data", chunk.to_base64(config));
            add_param(&mut params, "segment_index", chunk_num.to_string());

            let req = auth::post(links::medias::UPLOAD, &self.token, Some(&params));
            Some(make_parsed_future(&self.handle, req))
        } else {
            None
        }
    }

    fn finalize(&self, media_id: u64) -> FutureResponse<Media> {
        let mut params = HashMap::new();

        add_param(&mut params, "command", "FINALIZE");
        add_param(&mut params, "media_id", media_id.to_string());

        let req = auth::post(links::medias::UPLOAD, &self.token, Some(&params));
        make_parsed_future(&self.handle, req)
    }
}

impl<'a> Future for UploadFuture<'a> {
    type Item = Response<Media>;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        use std::mem::replace;

        match replace(&mut self.status, UploadInner::Invalid) {
            UploadInner::WaitingForInit(mut init) => {
                match init.poll() {
                    Ok(Async::NotReady) => {
                        self.status = UploadInner::WaitingForInit(init);
                        Ok(Async::NotReady)
                    },
                    Ok(Async::Ready(media)) => {
                        let id = media.id;
                        let loader = self.append(0, id).unwrap();
                        self.status = UploadInner::UploadingChunk(id, 0, loader);
                        self.poll()
                    },
                    Err(e) => Err(e),
                }
            },
            UploadInner::UploadingChunk(id, chunk_idx, mut upload) => {
                match upload.poll() {
                    Ok(Async::NotReady) => {
                        self.status = UploadInner::UploadingChunk(id, chunk_idx, upload);
                        Ok(Async::NotReady)
                    },
                    Ok(Async::Ready(_)) => {
                        let chunk_idx = chunk_idx + 1;
                        if let Some(upload) = self.append(chunk_idx, id) {
                            self.status = UploadInner::UploadingChunk(id, chunk_idx, upload);
                        } else {
                            let loader = self.finalize(id);
                            self.status = UploadInner::Finalizing(loader);
                        }

                        self.poll()
                    },
                    Err(e) => Err(e),
                }
            },
            UploadInner::Finalizing(mut finalize) => {
                match finalize.poll() {
                    Ok(Async::NotReady) => {
                        self.status = UploadInner::Finalizing(finalize);
                        Ok(Async::NotReady)
                    },
                    status => status,
                }
            },
            UploadInner::Invalid => Err(error::Error::FutureAlreadyCompleted),
        }
    }
}

/// Requests the current status of the given media upload.
///
/// Since videos and gifs are asynchronously processed, the media may not be ready even though the
/// `UploadFuture` has successfully resolved. To check in on the status of an upload, use this call
/// and inspect the `progress` of the returned `Media`. If it states `Some(Success)`, then the
/// media is ready to be attached to a new tweet.
pub fn status(id: u64, token: &auth::Token, handle: &Handle) -> FutureResponse<Media> {
    let mut params = HashMap::new();

    add_param(&mut params, "command", "STATUS");
    add_param(&mut params, "media_id", id.to_string());

    let req = auth::get(links::medias::UPLOAD, token, Some(&params));
    make_parsed_future(handle, req)
}

#[cfg(test)]
mod tests {
    use common::FromJson;

    use super::Media;

    use std::fs::File;
    use std::io::Read;

    fn load_media(path: &str) -> Media {
        let mut file = File::open(path).unwrap();
        let mut content = String::new();
        file.read_to_string(&mut content).unwrap();
        Media::from_str(&content).unwrap()
    }

    #[test]
    fn parse_media() {
        let media = load_media("src/media/media.json");

        assert_eq!(media.id, 710511363345354753);
        assert_eq!(media.expires_after, 86400);
    }

    #[test]
    fn parse_media_pending() {
        let media = load_media("src/media/media_pending.json");

        assert_eq!(media.id, 13);
        assert_eq!(media.expires_after, 86400);
        assert!(media.progress.is_some());

        match media.progress {
            Some(super::ProgressInfo::Pending(5)) => (),
            other => assert!(false, format!("Unexpected value of progress={:?}", other))
        }
    }

    #[test]
    fn parse_media_in_progress() {
        let media = load_media("src/media/media_in_progress.json");

        assert_eq!(media.id, 13);
        assert_eq!(media.expires_after, 3595);
        assert!(media.progress.is_some());

        match media.progress {
            Some(super::ProgressInfo::InProgress(10)) => (),
            other => assert!(false, format!("Unexpected value of progress={:?}", other))
        }
    }

    #[test]
    fn parse_media_fail() {
        let media = load_media("src/media/media_fail.json");

        assert_eq!(media.id, 710511363345354753);
        assert_eq!(media.expires_after, 0);
        assert!(media.progress.is_some());

        match media.progress {
            Some(super::ProgressInfo::Failed(error)) => assert_eq!(error, "InvalidMedia: Unsupported video format"),
            other => assert!(false, format!("Unexpected value of progress={:?}", other))
        }
    }
}
<commit_msg>parse/return the full processing error<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//!Media module
//!
//!Provides functionality to upload images, GIFs and videos using Twitter Media API

use std::borrow::Cow;
use std::collections::HashMap;

use futures::{Future, Async, Poll};
use rustc_serialize::{json, base64};
use rustc_serialize::base64::{ToBase64};

use common::*;
use error;
use error::Error::InvalidResponse;
use links;
use auth;

use mime;

/// A collection of convenience functions that return media types accepted by Twitter.
pub mod media_types {
    use mime::{self, Mime};

    /// PNG images.
    pub fn image_png() -> Mime {
        mime::IMAGE_PNG
    }

    /// JPG images.
    pub fn image_jpg() -> Mime {
        mime::IMAGE_JPEG
    }

    /// WEBP images.
    pub fn image_webp() -> Mime {
        "image/webp".parse().unwrap()
    }

    /// GIF images, both animated and static.
    pub fn image_gif() -> Mime {
        mime::IMAGE_GIF
    }

    /// MP4 videos.
    pub fn video_mp4() -> Mime {
        "video/mp4".parse().unwrap()
    }
}

///Media's upload progressing info.
#[derive(Debug)]
pub enum ProgressInfo {
    ///Video is pending for processing. Contains number of seconds after which to check.
    Pending(u64),
    ///Video is beeing processed. Contains number of seconds after which to check.
    InProgress(u64),
    ///Video's processing failed. Contains reason.
    Failed(MediaError),
    ///Video's processing is finished. Media can be used in other API calls.
    Success
}

/// Represents an error that can occur during media processing.
#[derive(Debug)]
pub struct MediaError {
    /// A numeric error code assigned to the error.
    pub code: i32,
    /// A short name given to the error.
    pub name: String,
    /// The full text of the error message.
    pub message: String,
}

impl FromJson for ProgressInfo {
    fn from_json(input: &json::Json) -> Result<Self, error::Error> {
        field_present!(input, state);
        let state: String = try!(field(input, "state"));

        match state.as_ref() {
            "pending" => {
                field_present!(input, check_after_secs);
                Ok(ProgressInfo::Pending(try!(field(input, "check_after_secs"))))
            },
            "in_progress" => {
                field_present!(input, check_after_secs);
                Ok(ProgressInfo::InProgress(try!(field(input, "check_after_secs"))))
            },
            "failed" => {
                field_present!(input, error);
                Ok(ProgressInfo::Failed(try!(field(input, "error"))))
            },
            "succeeded" => Ok(ProgressInfo::Success),
            state => Err(InvalidResponse("Unexpected progress info state", Some(state.to_string())))

        }
    }
}

impl FromJson for MediaError {
    fn from_json(input: &json::Json) -> Result<Self, error::Error> {
        field_present!(input, code);
        field_present!(input, name);
        field_present!(input, message);

        Ok(MediaError {
            code: try!(field(input, "code")),
            name: try!(field(input, "name")),
            message: try!(field(input, "message")),
        })
    }
}

///Represents media file that is uploaded on twitter.
pub struct Media {
    ///ID that can be used in API calls (e.g. attach to tweet).
    pub id: u64,
    ///Number of second the media can be used in other API calls.
    pub expires_after: u64,
    ///Progress information. If present determines whether Media can be used.
    pub progress: Option<ProgressInfo>
}

impl FromJson for Media {
    fn from_json(input: &json::Json) -> Result<Self, error::Error> {
        if !input.is_object() {
            return Err(InvalidResponse("Tweet received json that wasn't an object", Some(input.to_string())));
        }

        field_present!(input, media_id);

        Ok(Media {
            id: try!(field(input, "media_id")),
            //We can miss this field on failed upload in which case 0 is pretty reasonable value.
            expires_after: field(input, "expires_after_secs").unwrap_or(0),
            progress: try!(field(input, "processing_info"))
        })
    }
}

///Uploads image using old twitter API.
///
///The image should be raw binary with content of file.
pub fn upload_image(image: &[u8], token: &auth::Token, handle: &Handle) -> FutureResponse<Media> {
    let mut params = HashMap::new();

    let config = base64::Config {
        char_set: base64::CharacterSet::Standard,
        newline: base64::Newline::LF,
        pad: true,
        line_length: None,
    };

    add_param(&mut params, "media_data", image.to_base64(config));

    let req = auth::post(links::medias::UPLOAD, token, Some(&params));
    make_parsed_future(handle, req)
}

/// Represents the kinda of media that Twitter will accept.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum MediaCategory {
    /// Static image. Four can be attached to a single tweet.
    Image,
    /// Animated GIF.
    Gif,
    /// Video.
    Video,
}

/// `Display` impl for `MediaCategory` so that `.to_string()` will return a string suitable for use
/// in an API call. This will turn the enum into `"tweet_image"`, `"tweet_gif"`, and
/// `"tweet_video"`.
impl ::std::fmt::Display for MediaCategory {
    fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match *self {
            MediaCategory::Image => write!(fmt, "tweet_image"),
            MediaCategory::Gif => write!(fmt, "tweet_gif"),
            MediaCategory::Video => write!(fmt, "tweet_video"),
        }
    }
}

/// A builder struct that allows you to build up parameters to a media upload before initiating it.
pub struct UploadBuilder<'a> {
    data: Cow<'a, [u8]>,
    media_type: mime::Mime,
    chunk_size: Option<usize>,
    category: Option<MediaCategory>,
}

impl<'a> UploadBuilder<'a> {
    /// Begins setting up a media upload call, with the given data and media type.
    ///
    /// For convenience functions to get known `media_type`s that Twitter will accept, see the
    /// [`media_types`] module.
    ///
    /// [`media_types`]: media_types/index.html
    pub fn new<V: Into<Cow<'a, [u8]>>>(data: V, media_type: mime::Mime) -> UploadBuilder<'a> {
        UploadBuilder {
            data: data.into(),
            media_type,
            chunk_size: None,
            category: None,
        }
    }

    /// Sets how many bytes to upload in one network call. By default this is set to 512 KiB.
    pub fn chunk_size(self, chunk_size: usize) -> Self {
        UploadBuilder {
            chunk_size: Some(chunk_size),
            ..self
        }
    }

    /// Sets the `media_category` sent to Twitter. When unset, it behaves as if you sent
    /// `MediaCategory::Image`.
    pub fn category(self, category: MediaCategory) -> Self {
        UploadBuilder {
            category: Some(category),
            ..self
        }
    }

    /// Collects the built-up parameters and begins the chunked upload.
    pub fn call(self, token: &auth::Token, handle: &Handle) -> UploadFuture<'a> {
        let mut params = HashMap::new();

        add_param(&mut params, "command", "INIT");
        add_param(&mut params, "total_bytes", self.data.len().to_string());
        add_param(&mut params, "media_type", self.media_type.to_string());

        if let Some(category) = self.category {
            add_param(&mut params, "media_category", category.to_string());
        }

        let req = auth::post(links::medias::UPLOAD, token, Some(&params));
        let loader = make_parsed_future(handle, req);
        UploadFuture {
            data: self.data,
            token: token.clone(),
            handle: handle.clone(),
            chunk_size: self.chunk_size.unwrap_or(1024 * 512), // 512 KiB default
            status: UploadInner::WaitingForInit(loader),
        }
    }
}

/// A `Future` that represents an in-progress media upload.
pub struct UploadFuture<'a> {
    data: Cow<'a, [u8]>,
    token: auth::Token,
    handle: Handle,
    chunk_size: usize,
    status: UploadInner,
}

/// The current status of an `UploadFuture`.
enum UploadInner {
    /// The `UploadFuture` is waiting to initialize the media upload session.
    WaitingForInit(FutureResponse<Media>),
    /// The `UploadFuture` is in the progress of uploading data.
    UploadingChunk(u64, usize, FutureResponse<()>),
    /// The `UploadFuture` is currently finalizing the media with Twitter.
    Finalizing(FutureResponse<Media>),
    /// The `UploadFuture` has completed, or has encountered an error.
    Invalid,
}

impl<'a> UploadFuture<'a> {
    fn get_chunk(&self, chunk_num: usize) -> Option<&[u8]> {
        let start = chunk_num * self.chunk_size;
        let end = (chunk_num + 1) * self.chunk_size;

        if start >= self.data.len() {
            None
        } else if end >= self.data.len() {
            Some(&self.data[start..])
        } else {
            Some(&self.data[start..end])
        }
    }

    fn append(&self, chunk_num: usize, media_id: u64) -> Option<FutureResponse<()>> {
        let mut chunk = self.get_chunk(chunk_num);
        if chunk.is_none() && chunk_num == 0 {
            chunk = Some(&[][..]);
        }

        if let Some(chunk) = chunk {
            let mut params = HashMap::new();

            let config = base64::Config {
                char_set: base64::CharacterSet::Standard,
                newline: base64::Newline::LF,
                pad: true,
                line_length: None,
            };

            add_param(&mut params, "command", "APPEND");
            add_param(&mut params, "media_id", media_id.to_string());
            add_param(&mut params, "media_data", chunk.to_base64(config));
            add_param(&mut params, "segment_index", chunk_num.to_string());

            let req = auth::post(links::medias::UPLOAD, &self.token, Some(&params));
            Some(make_parsed_future(&self.handle, req))
        } else {
            None
        }
    }

    fn finalize(&self, media_id: u64) -> FutureResponse<Media> {
        let mut params = HashMap::new();

        add_param(&mut params, "command", "FINALIZE");
        add_param(&mut params, "media_id", media_id.to_string());

        let req = auth::post(links::medias::UPLOAD, &self.token, Some(&params));
        make_parsed_future(&self.handle, req)
    }
}

impl<'a> Future for UploadFuture<'a> {
    type Item = Response<Media>;
    type Error = error::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        use std::mem::replace;

        match replace(&mut self.status, UploadInner::Invalid) {
            UploadInner::WaitingForInit(mut init) => {
                match init.poll() {
                    Ok(Async::NotReady) => {
                        self.status = UploadInner::WaitingForInit(init);
                        Ok(Async::NotReady)
                    },
                    Ok(Async::Ready(media)) => {
                        let id = media.id;
                        let loader = self.append(0, id).unwrap();
                        self.status = UploadInner::UploadingChunk(id, 0, loader);
                        self.poll()
                    },
                    Err(e) => Err(e),
                }
            },
            UploadInner::UploadingChunk(id, chunk_idx, mut upload) => {
                match upload.poll() {
                    Ok(Async::NotReady) => {
                        self.status = UploadInner::UploadingChunk(id, chunk_idx, upload);
                        Ok(Async::NotReady)
                    },
                    Ok(Async::Ready(_)) => {
                        let chunk_idx = chunk_idx + 1;
                        if let Some(upload) = self.append(chunk_idx, id) {
                            self.status = UploadInner::UploadingChunk(id, chunk_idx, upload);
                        } else {
                            let loader = self.finalize(id);
                            self.status = UploadInner::Finalizing(loader);
                        }

                        self.poll()
                    },
                    Err(e) => Err(e),
                }
            },
            UploadInner::Finalizing(mut finalize) => {
                match finalize.poll() {
                    Ok(Async::NotReady) => {
                        self.status = UploadInner::Finalizing(finalize);
                        Ok(Async::NotReady)
                    },
                    status => status,
                }
            },
            //TODO: you know what, weave status in here too, call in tokio-timer or something to
            //sleep and make a new error variant to include media processing errors
            UploadInner::Invalid => Err(error::Error::FutureAlreadyCompleted),
        }
    }
}

/// Requests the current status of the given media upload.
///
/// Since videos and gifs are asynchronously processed, the media may not be ready even though the
/// `UploadFuture` has successfully resolved. To check in on the status of an upload, use this call
/// and inspect the `progress` of the returned `Media`. If it states `Some(Success)`, then the
/// media is ready to be attached to a new tweet.
pub fn status(id: u64, token: &auth::Token, handle: &Handle) -> FutureResponse<Media> {
    let mut params = HashMap::new();

    add_param(&mut params, "command", "STATUS");
    add_param(&mut params, "media_id", id.to_string());

    let req = auth::get(links::medias::UPLOAD, token, Some(&params));
    make_parsed_future(handle, req)
}

#[cfg(test)]
mod tests {
    use common::FromJson;

    use super::Media;

    use std::fs::File;
    use std::io::Read;

    fn load_media(path: &str) -> Media {
        let mut file = File::open(path).unwrap();
        let mut content = String::new();
        file.read_to_string(&mut content).unwrap();
        Media::from_str(&content).unwrap()
    }

    #[test]
    fn parse_media() {
        let media = load_media("src/media/media.json");

        assert_eq!(media.id, 710511363345354753);
        assert_eq!(media.expires_after, 86400);
    }

    #[test]
    fn parse_media_pending() {
        let media = load_media("src/media/media_pending.json");

        assert_eq!(media.id, 13);
        assert_eq!(media.expires_after, 86400);
        assert!(media.progress.is_some());

        match media.progress {
            Some(super::ProgressInfo::Pending(5)) => (),
            other => assert!(false, format!("Unexpected value of progress={:?}", other))
        }
    }

    #[test]
    fn parse_media_in_progress() {
        let media = load_media("src/media/media_in_progress.json");

        assert_eq!(media.id, 13);
        assert_eq!(media.expires_after, 3595);
        assert!(media.progress.is_some());

        match media.progress {
            Some(super::ProgressInfo::InProgress(10)) => (),
            other => assert!(false, format!("Unexpected value of progress={:?}", other))
        }
    }

    #[test]
    fn parse_media_fail() {
        let media = load_media("src/media/media_fail.json");

        assert_eq!(media.id, 710511363345354753);
        assert_eq!(media.expires_after, 0);
        assert!(media.progress.is_some());

        match media.progress {
            Some(super::ProgressInfo::Failed(error)) => assert_eq!(error, "InvalidMedia: Unsupported video format"),
            other => assert!(false, format!("Unexpected value of progress={:?}", other))
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate rayon;
extern crate rusoto_codegen;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate toml;

mod cargo;

use std::collections::BTreeMap;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write, BufWriter};
use std::path::Path;

use rayon::prelude::*;

use clap::{Arg, ArgGroup, App};

use rusoto_codegen::botocore::Service;

#[derive(Deserialize)]
struct ServiceConfig {
    pub version: String
}

fn get_dependencies(service: &Service, core_version: &str) -> BTreeMap<String, cargo::Dependency> {
    let mut dependencies = BTreeMap::new();

    dependencies.insert("hyper".to_owned(), cargo::Dependency::Simple("0.10.0".into()));
    dependencies.insert("rusoto_core".to_owned(), cargo::Dependency::Extended {
        path: Some("../../core".into()),
        version: Some(core_version.to_owned()),
        optional: None,
        default_features: None,
        features: None
    });

    match &service.metadata.protocol[..] {
        "json" => {
            dependencies.insert("serde".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_derive".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_json".to_owned(), cargo::Dependency::Simple("1.0.1".into()));
        },
        "query" | "ec2" => {
            dependencies.insert("xml-rs".to_owned(), cargo::Dependency::Simple("0.3".into()));
        },
        "rest-json" => {
            dependencies.insert("log".to_owned(), cargo::Dependency::Simple("0.3.6".into()));
            dependencies.insert("serde".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_derive".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_json".to_owned(), cargo::Dependency::Simple("1.0.1".into()));
        },
        "rest-xml" => {
            dependencies.insert("xml-rs".to_owned(), cargo::Dependency::Simple("0.3".into()));
        },
        protocol => panic!("Unknown protocol {}", protocol),
    }

    match &service.metadata.endpoint_prefix[..] {
        "s3" => {
            dependencies.insert("rustc-serialize".to_owned(), cargo::Dependency::Simple("0.3.19".into()));
            dependencies.insert("md5".to_owned(), cargo::Dependency::Simple("0.3.2".into()));
        },
        "sts" => {
            dependencies.insert("chrono".to_owned(), cargo::Dependency::Simple("0.2.25".into()));
        }
        _ => {}
    };

    dependencies
}

fn main() {
    let matches = App::new("Rusoto Service Crate Generator")
        .version(crate_version!())
        .author(crate_authors!())
        .about(crate_description!())
        .arg(Arg::with_name("services_list")
            .long("services")
            .use_delimiter(true)
            .takes_value(true))
        .arg(Arg::with_name("services_config")
            .long("config")
            .short("c")
            .takes_value(true))
        .group(ArgGroup::with_name("services")
            .args(&["services_list", "services_config"])
            .required(true))
        .arg(Arg::with_name("out_dir")
            .long("outdir")
            .short("o")
            .takes_value(true)
            .required(true))
        .get_matches();

    let core_manifest_raw = OpenOptions::new()
        .read(true)
        .open("../rusoto/core/Cargo.toml")
        .and_then(|mut f| {
            let mut buf = String::new();
            f.read_to_string(&mut buf).map(|_| buf)
        })
        .expect("Unable to read core crate's Cargo.toml");

    let core_manifest: cargo::Manifest = toml::from_str(&core_manifest_raw).expect("Unable to parse core crate's Cargo.toml");

    let version = &core_manifest.package.version;

    let out_dir = Path::new(matches.value_of("out_dir").unwrap());

    if !out_dir.exists() {
        fs::create_dir(out_dir).expect("Unable to create output directory");
    }

    let service_versions: Vec<(String, String)> = {
        if let Some(services) = matches.values_of("services_list") {
            services.map(|s| {
                let mut service_parts = s.splitn(2, "@");
                let name = service_parts.next().expect(&format!("Invalid service value {}. Must be in format name@version.", s)).to_owned();
                let version = service_parts.next().expect(&format!("Invalid service value {}. Must be in format name@version.", s)).to_owned();
                (name, version)
            }).collect()
        } else if let Some(services_config) = matches.value_of("services_config") {
            let contents = File::open(services_config).and_then(|mut f| {
                let mut contents = String::new();
                f.read_to_string(&mut contents).map(|_| contents)
            }).expect("Unable to read services configuration file.");

            let parsed: BTreeMap<String, ServiceConfig> = serde_json::from_str(&contents).expect("Unable to parse services configuration file.");
            parsed.iter().map(|(k, v)| (k.clone(), v.version.clone())).collect()
        } else {
            unreachable!("Arg parsing failed. Did not have services config or services list.");
        }
    };

    service_versions.par_iter().for_each(|&(ref name, ref service_version)| {
        let service = Service::load(name, service_version);

        if let Err(ref e) = service {
            println!("Failed to load service {}: {}", name, e);
            return;
        }

        let service = service.as_ref().unwrap();

        let crate_dir = out_dir.join(&name);
        let crate_name = format!("rusoto_{}", &name);

        println!("Generating crate for {} @ {}...", service.metadata.service_full_name, service.metadata.api_version);

        if !crate_dir.exists() {
            fs::create_dir(&crate_dir).expect(&format!("Unable to create directory at {}", crate_dir.display()));
        }

        let service_dependencies = get_dependencies(&service, &version);

        let extern_crates = service_dependencies.iter().map(|(k, _)| {
            if k == "xml-rs" {
                return "extern crate xml;".into();
            }
            let safe_name = k.replace("-", "_");
            let use_macro = k == "serde_derive" || k == "log" || k == "lazy_static";
            if use_macro {
                return format!("#[macro_use]\nextern crate {};", safe_name);
            }
            format!("extern crate {};", safe_name)
        }).collect::<Vec<String>>().join("\n");

        let mut cargo_manifest = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(crate_dir.join("Cargo.toml"))
            .expect("Unable to write Cargo.toml");

        let manifest = cargo::Manifest {
            package: cargo::Metadata {
                authors: Some(vec![
                    "Anthony DiMarco <ocramida@gmail.com>".into(),
                    "Jimmy Cuadra <jimmy@jimmycuadra.com>".into(),
                    "Matthew Mayer <matthewkmayer@gmail.com>".into(),
                    "Nikita Pekin <contact@nikitapek.in>".into()
                ]),
                description: Some(format!("AWS SDK for Rust - {} @ {}", &service.metadata.service_full_name, &service.metadata.api_version)),
                documentation: Some("http://rusoto.github.io/rusoto/rusoto/index.html".into()),
                keywords: Some(vec!["AWS".into(), "Amazon".into(), name.clone(), service.metadata.service_full_name.clone(), service.metadata.endpoint_prefix.clone()]),
                license: Some("MIT".into()),
                name: crate_name.clone(),
                readme: None,
                repository: Some("https://github.com/rusoto/rusoto".into()),
                version: version.clone(),
                homepage: Some("https://www.rusoto.org/".into()),
                ..cargo::Metadata::default()
            },
            dependencies: service_dependencies,
            dev_dependencies: vec![
                ("rusoto_mock".to_owned(), cargo::Dependency::Extended {
                    path: Some("../../mock".into()),
                    version: None,
                    optional: None,
                    default_features: None,
                    features: None
                })
            ].into_iter().collect(),
            ..cargo::Manifest::default()
        };

        cargo_manifest.write_all(toml::to_string(&manifest).unwrap().as_bytes()).unwrap();

        {
            let src_dir = crate_dir.join("src");

            if !src_dir.exists() {
                fs::create_dir(&src_dir).expect(&format!("Unable to create directory at {}", src_dir.display()));
            }

            let lib_file_path = src_dir.join("lib.rs");

            let lib_file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&lib_file_path)
                .expect("Unable to write lib.rs");

            let mut lib_writer = BufWriter::new(lib_file);

            lib_writer.write_all(extern_crates.as_bytes()).unwrap();
            lib_writer.write_all(b"\n\n").unwrap();
            lib_writer.write_all(b"mod generated;\npub use generated::*;\n").unwrap();
            lib_writer.write_all(b"mod custom;\npub use custom::*;\n").unwrap();

            let gen_file_path = src_dir.join("generated.rs");

            let gen_file = OpenOptions::new()
                .create(true)
                .write(true)
                .truncate(true)
                .open(&gen_file_path)
                .expect("Unable to write generated.rs");

            let mut gen_writer = BufWriter::new(gen_file);

            rusoto_codegen::generator::generate_source(&service, &mut gen_writer).unwrap();

            let custom_dir_path = src_dir.join("custom");

            if !custom_dir_path.exists() {
                fs::create_dir(&custom_dir_path).expect(&format!("Unable to create directory at {}", custom_dir_path.display()));
            }

            let custom_mod_file_path = custom_dir_path.join("mod.rs");

            if !custom_mod_file_path.exists() {
                OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&custom_mod_file_path)
                    .expect("Unable to write mod.rs");
            }
        }

        {
            let test_resources_dir = crate_dir.join("test_resources");

            if !test_resources_dir.exists() {
                fs::create_dir(&test_resources_dir).expect(&format!("Unable to create directory at {}", test_resources_dir.display()));
            }

            let test_resources = rusoto_codegen::generator::tests::find_responses_for_service(&service);
            if !test_resources.is_empty() {
                let generated_test_resources_dir = test_resources_dir.join("generated");

                if !generated_test_resources_dir.exists() {
                    fs::create_dir(&generated_test_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
                }

                for resource in test_resources {
                    fs::copy(resource.full_path, generated_test_resources_dir.join(&resource.file_name)).expect("Failed to copy test resource file");
                    // TODO: remove this. added to migrate
                    if test_resources_dir.join(&resource.file_name).exists() {
                        fs::remove_file(test_resources_dir.join(&resource.file_name)).unwrap();
                    }
                }
            }
        }
    });
}
<commit_msg>Remove TODO that is done<commit_after>#[macro_use]
extern crate clap;
extern crate rayon;
extern crate rusoto_codegen;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate toml;

mod cargo;

use std::collections::BTreeMap;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write, BufWriter};
use std::path::Path;

use rayon::prelude::*;

use clap::{Arg, ArgGroup, App};

use rusoto_codegen::botocore::Service;

#[derive(Deserialize)]
struct ServiceConfig {
    pub version: String
}

fn get_dependencies(service: &Service, core_version: &str) -> BTreeMap<String, cargo::Dependency> {
    let mut dependencies = BTreeMap::new();

    dependencies.insert("hyper".to_owned(), cargo::Dependency::Simple("0.10.0".into()));
    dependencies.insert("rusoto_core".to_owned(), cargo::Dependency::Extended {
        path: Some("../../core".into()),
        version: Some(core_version.to_owned()),
        optional: None,
        default_features: None,
        features: None
    });

    match &service.metadata.protocol[..] {
        "json" => {
            dependencies.insert("serde".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_derive".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_json".to_owned(), cargo::Dependency::Simple("1.0.1".into()));
        },
        "query" | "ec2" => {
            dependencies.insert("xml-rs".to_owned(), cargo::Dependency::Simple("0.3".into()));
        },
        "rest-json" => {
            dependencies.insert("log".to_owned(), cargo::Dependency::Simple("0.3.6".into()));
            dependencies.insert("serde".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_derive".to_owned(), cargo::Dependency::Simple("1.0.2".into()));
            dependencies.insert("serde_json".to_owned(), cargo::Dependency::Simple("1.0.1".into()));
        },
        "rest-xml" => {
            dependencies.insert("xml-rs".to_owned(), cargo::Dependency::Simple("0.3".into()));
        },
        protocol => panic!("Unknown protocol {}", protocol),
    }

    match &service.metadata.endpoint_prefix[..] {
        "s3" => {
            dependencies.insert("rustc-serialize".to_owned(), cargo::Dependency::Simple("0.3.19".into()));
            dependencies.insert("md5".to_owned(), cargo::Dependency::Simple("0.3.2".into()));
        },
        "sts" => {
            dependencies.insert("chrono".to_owned(), cargo::Dependency::Simple("0.2.25".into()));
        }
        _ => {}
    };

    dependencies
}

fn main() {
    let matches = App::new("Rusoto Service Crate Generator")
        .version(crate_version!())
        .author(crate_authors!())
        .about(crate_description!())
        .arg(Arg::with_name("services_list")
            .long("services")
            .use_delimiter(true)
            .takes_value(true))
        .arg(Arg::with_name("services_config")
            .long("config")
            .short("c")
            .takes_value(true))
        .group(ArgGroup::with_name("services")
            .args(&["services_list", "services_config"])
            .required(true))
        .arg(Arg::with_name("out_dir")
            .long("outdir")
            .short("o")
            .takes_value(true)
            .required(true))
        .get_matches();

    let core_manifest_raw = OpenOptions::new()
        .read(true)
        .open("../rusoto/core/Cargo.toml")
        .and_then(|mut f| {
            let mut buf = String::new();
            f.read_to_string(&mut buf).map(|_| buf)
        })
        .expect("Unable to read core crate's Cargo.toml");

    let core_manifest: cargo::Manifest = toml::from_str(&core_manifest_raw).expect("Unable to parse core crate's Cargo.toml");

    let version = &core_manifest.package.version;

    let out_dir = Path::new(matches.value_of("out_dir").unwrap());

    if !out_dir.exists() {
        fs::create_dir(out_dir).expect("Unable to create output directory");
    }

    let service_versions: Vec<(String, String)> = {
        if let Some(services) = matches.values_of("services_list") {
            services.map(|s| {
                let mut service_parts = s.splitn(2, "@");
                let name = service_parts.next().expect(&format!("Invalid service value {}. Must be in format name@version.", s)).to_owned();
                let version = service_parts.next().expect(&format!("Invalid service value {}. Must be in format name@version.", s)).to_owned();
                (name, version)
            }).collect()
        } else if let Some(services_config) = matches.value_of("services_config") {
            let contents = File::open(services_config).and_then(|mut f| {
                let mut contents = String::new();
                f.read_to_string(&mut contents).map(|_| contents)
            }).expect("Unable to read services configuration file.");

            let parsed: BTreeMap<String, ServiceConfig> = serde_json::from_str(&contents).expect("Unable to parse services configuration file.");
            parsed.iter().map(|(k, v)| (k.clone(), v.version.clone())).collect()
        } else {
            unreachable!("Arg parsing failed. Did not have services config or services list.");
        }
    };

    service_versions.par_iter().for_each(|&(ref name, ref service_version)| {
        let service = Service::load(name, service_version);

        if let Err(ref e) = service {
            println!("Failed to load service {}: {}", name, e);
            return;
        }

        let service = service.as_ref().unwrap();

        let crate_dir = out_dir.join(&name);
        let crate_name = format!("rusoto_{}", &name);

        println!("Generating crate for {} @ {}...", service.metadata.service_full_name, service.metadata.api_version);

        if !crate_dir.exists() {
            fs::create_dir(&crate_dir).expect(&format!("Unable to create directory at {}", crate_dir.display()));
        }

        let service_dependencies = get_dependencies(&service, &version);

        let extern_crates = service_dependencies.iter().map(|(k, _)| {
            if k == "xml-rs" {
                return "extern crate xml;".into();
            }
            let safe_name = k.replace("-", "_");
            let use_macro = k == "serde_derive" || k == "log" || k == "lazy_static";
            if use_macro {
                return format!("#[macro_use]\nextern crate {};", safe_name);
            }
            format!("extern crate {};", safe_name)
        }).collect::<Vec<String>>().join("\n");

        let mut cargo_manifest = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(crate_dir.join("Cargo.toml"))
            .expect("Unable to write Cargo.toml");

        let manifest = cargo::Manifest {
            package: cargo::Metadata {
                authors: Some(vec![
                    "Anthony DiMarco <ocramida@gmail.com>".into(),
                    "Jimmy Cuadra <jimmy@jimmycuadra.com>".into(),
                    "Matthew Mayer <matthewkmayer@gmail.com>".into(),
                    "Nikita Pekin <contact@nikitapek.in>".into()
                ]),
                description: Some(format!("AWS SDK for Rust - {} @ {}", &service.metadata.service_full_name, &service.metadata.api_version)),
                documentation: Some("http://rusoto.github.io/rusoto/rusoto/index.html".into()),
                keywords: Some(vec!["AWS".into(), "Amazon".into(), name.clone(), service.metadata.service_full_name.clone(), service.metadata.endpoint_prefix.clone()]),
                license: Some("MIT".into()),
                name: crate_name.clone(),
                readme: None,
                repository: Some("https://github.com/rusoto/rusoto".into()),
                version: version.clone(),
                homepage: Some("https://www.rusoto.org/".into()),
                ..cargo::Metadata::default()
            },
            dependencies: service_dependencies,
            dev_dependencies: vec![
                ("rusoto_mock".to_owned(), cargo::Dependency::Extended {
                    path: Some("../../mock".into()),
                    version: None,
                    optional: None,
                    default_features: None,
                    features: None
                })
            ].into_iter().collect(),
            ..cargo::Manifest::default()
        };

        cargo_manifest.write_all(toml::to_string(&manifest).unwrap().as_bytes()).unwrap();

        {
            let src_dir = crate_dir.join("src");

            if !src_dir.exists() {
                fs::create_dir(&src_dir).expect(&format!("Unable to create directory at {}", src_dir.display()));
            }

            let lib_file_path = src_dir.join("lib.rs");

            let lib_file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&lib_file_path)
                .expect("Unable to write lib.rs");

            let mut lib_writer = BufWriter::new(lib_file);

            lib_writer.write_all(extern_crates.as_bytes()).unwrap();
            lib_writer.write_all(b"\n\n").unwrap();
            lib_writer.write_all(b"mod generated;\npub use generated::*;\n").unwrap();
            lib_writer.write_all(b"mod custom;\npub use custom::*;\n").unwrap();

            let gen_file_path = src_dir.join("generated.rs");

            let gen_file = OpenOptions::new()
                .create(true)
                .write(true)
                .truncate(true)
                .open(&gen_file_path)
                .expect("Unable to write generated.rs");

            let mut gen_writer = BufWriter::new(gen_file);

            rusoto_codegen::generator::generate_source(&service, &mut gen_writer).unwrap();

            let custom_dir_path = src_dir.join("custom");

            if !custom_dir_path.exists() {
                fs::create_dir(&custom_dir_path).expect(&format!("Unable to create directory at {}", custom_dir_path.display()));
            }

            let custom_mod_file_path = custom_dir_path.join("mod.rs");

            if !custom_mod_file_path.exists() {
                OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&custom_mod_file_path)
                    .expect("Unable to write mod.rs");
            }
        }

        {
            let test_resources_dir = crate_dir.join("test_resources");

            if !test_resources_dir.exists() {
                fs::create_dir(&test_resources_dir).expect(&format!("Unable to create directory at {}", test_resources_dir.display()));
            }

            let test_resources = rusoto_codegen::generator::tests::find_responses_for_service(&service);
            if !test_resources.is_empty() {
                let generated_test_resources_dir = test_resources_dir.join("generated");

                if !generated_test_resources_dir.exists() {
                    fs::create_dir(&generated_test_resources_dir).expect(&format!("Unable to create directory at {}", generated_test_resources_dir.display()));
                }

                for resource in test_resources {
                    fs::copy(resource.full_path, generated_test_resources_dir.join(&resource.file_name)).expect("Failed to copy test resource file");
                }
            }
        }
    });
}
<|endoftext|>


<commit_before>// ignore-windows: Concurrency on Windows is not supported yet.
// compile-flags: -Zmiri-disable-isolation -Zmiri-strict-provenance

use std::sync::{Arc, Barrier, Condvar, Mutex, Once, RwLock};
use std::thread;
use std::time::{Duration, Instant};

// Check if Rust barriers are working.

/// This test is taken from the Rust documentation.
fn check_barriers() {
    let mut handles = Vec::with_capacity(10);
    let barrier = Arc::new(Barrier::new(10));
    for _ in 0..10 {
        let c = barrier.clone();
        // The same messages will be printed together.
        // You will NOT see any interleaving.
        handles.push(thread::spawn(move || {
            println!("before wait");
            c.wait();
            println!("after wait");
        }));
    }
    // Wait for other threads to finish.
    for handle in handles {
        handle.join().unwrap();
    }
}

// Check if Rust conditional variables are working.

/// The test taken from the Rust documentation.
fn check_conditional_variables_notify_one() {
    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let pair2 = pair.clone();

    // Spawn a new thread.
    thread::spawn(move || {
        thread::yield_now();
        let (lock, cvar) = &*pair2;
        let mut started = lock.lock().unwrap();
        *started = true;
        // We notify the condvar that the value has changed.
        cvar.notify_one();
    });

    // Wait for the thread to fully start up.
    let (lock, cvar) = &*pair;
    let mut started = lock.lock().unwrap();
    while !*started {
        started = cvar.wait(started).unwrap();
    }
}

fn check_conditional_variables_notify_all() {
    let pair = Arc::new(((Mutex::new(())), Condvar::new()));

    // Spawn threads and block them on the conditional variable.
    let handles: Vec<_> = (0..5)
        .map(|_| {
            let pair2 = pair.clone();
            thread::spawn(move || {
                let (lock, cvar) = &*pair2;
                let guard = lock.lock().unwrap();
                // Block waiting on the conditional variable.
                let _ = cvar.wait(guard).unwrap();
            })
        })
        .inspect(|_| {
            thread::yield_now();
            thread::yield_now();
        })
        .collect();

    let (_, cvar) = &*pair;
    // Unblock all threads.
    cvar.notify_all();

    for handle in handles {
        handle.join().unwrap();
    }
}

/// Test that waiting on a conditional variable with a timeout does not
/// deadlock.
fn check_conditional_variables_timed_wait_timeout() {
    let lock = Mutex::new(());
    let cvar = Condvar::new();
    let guard = lock.lock().unwrap();
    let now = Instant::now();
    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(100)).unwrap();
    assert!(timeout.timed_out());
    let elapsed_time = now.elapsed().as_millis();
    assert!(100 <= elapsed_time && elapsed_time <= 500);
}

/// Test that signaling a conditional variable when waiting with a timeout works
/// as expected.
fn check_conditional_variables_timed_wait_notimeout() {
    let pair = Arc::new((Mutex::new(()), Condvar::new()));
    let pair2 = pair.clone();

    let (lock, cvar) = &*pair;
    let guard = lock.lock().unwrap();

    let handle = thread::spawn(move || {
        let (_lock, cvar) = &*pair2;
        cvar.notify_one();
    });

    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(500)).unwrap();
    assert!(!timeout.timed_out());
    handle.join().unwrap();
}

// Check if locks are working.

fn check_mutex() {
    let data = Arc::new(Mutex::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.lock().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_lock().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_write() {
    let data = Arc::new(RwLock::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.write().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_write().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_read_no_deadlock() {
    let l1 = Arc::new(RwLock::new(0));
    let l2 = Arc::new(RwLock::new(0));

    let l1_copy = Arc::clone(&l1);
    let l2_copy = Arc::clone(&l2);
    let _guard1 = l1.read().unwrap();
    let handle = thread::spawn(move || {
        let _guard2 = l2_copy.read().unwrap();
        thread::yield_now();
        let _guard1 = l1_copy.read().unwrap();
    });
    thread::yield_now();
    let _guard2 = l2.read().unwrap();
    handle.join().unwrap();
}

// Check if Rust once statics are working.

static mut VAL: usize = 0;
static INIT: Once = Once::new();

fn get_cached_val() -> usize {
    unsafe {
        INIT.call_once(|| {
            VAL = expensive_computation();
        });
        VAL
    }
}

fn expensive_computation() -> usize {
    let mut i = 1;
    let mut c = 1;
    while i < 1000 {
        i *= c;
        c += 1;
    }
    i
}

/// The test taken from the Rust documentation.
fn check_once() {
    let handles: Vec<_> = (0..10)
        .map(|_| {
            thread::spawn(|| {
                thread::yield_now();
                let val = get_cached_val();
                assert_eq!(val, 5040);
            })
        })
        .collect();
    for handle in handles {
        handle.join().unwrap();
    }
}

fn check_rwlock_unlock_bug1() {
    // There was a bug where when un-read-locking an rwlock that still has other
    // readers waiting, we'd accidentally also let a writer in.
    // That caused an ICE.
    let l = Arc::new(RwLock::new(0));

    let r1 = l.read().unwrap();
    let r2 = l.read().unwrap();

    // Make a waiting writer.
    let l2 = l.clone();
    thread::spawn(move || {
        let mut w = l2.write().unwrap();
        *w += 1;
    });
    thread::yield_now();

    drop(r1);
    assert_eq!(*r2, 0);
    thread::yield_now();
    thread::yield_now();
    thread::yield_now();
    assert_eq!(*r2, 0);
    drop(r2);
}

fn check_rwlock_unlock_bug2() {
    // There was a bug where when un-read-locking an rwlock by letting the last reader leaver,
    // we'd forget to wake up a writer.
    // That meant the writer thread could never run again.
    let l = Arc::new(RwLock::new(0));

    let r = l.read().unwrap();

    // Make a waiting writer.
    let l2 = l.clone();
    let h = thread::spawn(move || {
        let _w = l2.write().unwrap();
    });
    thread::yield_now();

    drop(r);
    h.join().unwrap();
}

fn park_timeout() {
    let start = Instant::now();

    thread::park_timeout(Duration::from_millis(200));
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn park_unpark() {
    let t1 = thread::current();
    let t2 = thread::spawn(move || {
        thread::park();
        thread::sleep(Duration::from_millis(200));
        t1.unpark();
    });

    let start = Instant::now();

    t2.thread().unpark();
    thread::park();
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn check_condvar() {
    let _ = std::sync::Condvar::new();
}

fn main() {
    check_barriers();
    check_conditional_variables_notify_one();
    check_conditional_variables_notify_all();
    check_conditional_variables_timed_wait_timeout();
    check_conditional_variables_timed_wait_notimeout();
    check_mutex();
    check_rwlock_write();
    check_rwlock_read_no_deadlock();
    check_once();
    check_rwlock_unlock_bug1();
    check_rwlock_unlock_bug2();
    park_timeout();
    park_unpark();
    check_condvar();
}
<commit_msg>Auto merge of #2063 - RalfJung:macos-timeout, r=RalfJung<commit_after>// ignore-windows: Concurrency on Windows is not supported yet.
// compile-flags: -Zmiri-disable-isolation -Zmiri-strict-provenance

use std::sync::{Arc, Barrier, Condvar, Mutex, Once, RwLock};
use std::thread;
use std::time::{Duration, Instant};

// Check if Rust barriers are working.

/// This test is taken from the Rust documentation.
fn check_barriers() {
    let mut handles = Vec::with_capacity(10);
    let barrier = Arc::new(Barrier::new(10));
    for _ in 0..10 {
        let c = barrier.clone();
        // The same messages will be printed together.
        // You will NOT see any interleaving.
        handles.push(thread::spawn(move || {
            println!("before wait");
            c.wait();
            println!("after wait");
        }));
    }
    // Wait for other threads to finish.
    for handle in handles {
        handle.join().unwrap();
    }
}

// Check if Rust conditional variables are working.

/// The test taken from the Rust documentation.
fn check_conditional_variables_notify_one() {
    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let pair2 = pair.clone();

    // Spawn a new thread.
    thread::spawn(move || {
        thread::yield_now();
        let (lock, cvar) = &*pair2;
        let mut started = lock.lock().unwrap();
        *started = true;
        // We notify the condvar that the value has changed.
        cvar.notify_one();
    });

    // Wait for the thread to fully start up.
    let (lock, cvar) = &*pair;
    let mut started = lock.lock().unwrap();
    while !*started {
        started = cvar.wait(started).unwrap();
    }
}

fn check_conditional_variables_notify_all() {
    let pair = Arc::new(((Mutex::new(())), Condvar::new()));

    // Spawn threads and block them on the conditional variable.
    let handles: Vec<_> = (0..5)
        .map(|_| {
            let pair2 = pair.clone();
            thread::spawn(move || {
                let (lock, cvar) = &*pair2;
                let guard = lock.lock().unwrap();
                // Block waiting on the conditional variable.
                let _ = cvar.wait(guard).unwrap();
            })
        })
        .inspect(|_| {
            thread::yield_now();
            thread::yield_now();
        })
        .collect();

    let (_, cvar) = &*pair;
    // Unblock all threads.
    cvar.notify_all();

    for handle in handles {
        handle.join().unwrap();
    }
}

/// Test that waiting on a conditional variable with a timeout does not
/// deadlock.
fn check_conditional_variables_timed_wait_timeout() {
    let lock = Mutex::new(());
    let cvar = Condvar::new();
    let guard = lock.lock().unwrap();
    let now = Instant::now();
    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(100)).unwrap();
    assert!(timeout.timed_out());
    let elapsed_time = now.elapsed().as_millis();
    assert!(100 <= elapsed_time && elapsed_time <= 800);
}

/// Test that signaling a conditional variable when waiting with a timeout works
/// as expected.
fn check_conditional_variables_timed_wait_notimeout() {
    let pair = Arc::new((Mutex::new(()), Condvar::new()));
    let pair2 = pair.clone();

    let (lock, cvar) = &*pair;
    let guard = lock.lock().unwrap();

    let handle = thread::spawn(move || {
        let (_lock, cvar) = &*pair2;
        cvar.notify_one();
    });

    let (_guard, timeout) = cvar.wait_timeout(guard, Duration::from_millis(500)).unwrap();
    assert!(!timeout.timed_out());
    handle.join().unwrap();
}

// Check if locks are working.

fn check_mutex() {
    let data = Arc::new(Mutex::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.lock().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_lock().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_write() {
    let data = Arc::new(RwLock::new(0));
    let mut threads = Vec::new();

    for _ in 0..3 {
        let data = Arc::clone(&data);
        let thread = thread::spawn(move || {
            let mut data = data.write().unwrap();
            thread::yield_now();
            *data += 1;
        });
        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    assert!(data.try_write().is_ok());

    let data = Arc::try_unwrap(data).unwrap().into_inner().unwrap();
    assert_eq!(data, 3);
}

fn check_rwlock_read_no_deadlock() {
    let l1 = Arc::new(RwLock::new(0));
    let l2 = Arc::new(RwLock::new(0));

    let l1_copy = Arc::clone(&l1);
    let l2_copy = Arc::clone(&l2);
    let _guard1 = l1.read().unwrap();
    let handle = thread::spawn(move || {
        let _guard2 = l2_copy.read().unwrap();
        thread::yield_now();
        let _guard1 = l1_copy.read().unwrap();
    });
    thread::yield_now();
    let _guard2 = l2.read().unwrap();
    handle.join().unwrap();
}

// Check if Rust once statics are working.

static mut VAL: usize = 0;
static INIT: Once = Once::new();

fn get_cached_val() -> usize {
    unsafe {
        INIT.call_once(|| {
            VAL = expensive_computation();
        });
        VAL
    }
}

fn expensive_computation() -> usize {
    let mut i = 1;
    let mut c = 1;
    while i < 1000 {
        i *= c;
        c += 1;
    }
    i
}

/// The test taken from the Rust documentation.
fn check_once() {
    let handles: Vec<_> = (0..10)
        .map(|_| {
            thread::spawn(|| {
                thread::yield_now();
                let val = get_cached_val();
                assert_eq!(val, 5040);
            })
        })
        .collect();
    for handle in handles {
        handle.join().unwrap();
    }
}

fn check_rwlock_unlock_bug1() {
    // There was a bug where when un-read-locking an rwlock that still has other
    // readers waiting, we'd accidentally also let a writer in.
    // That caused an ICE.
    let l = Arc::new(RwLock::new(0));

    let r1 = l.read().unwrap();
    let r2 = l.read().unwrap();

    // Make a waiting writer.
    let l2 = l.clone();
    thread::spawn(move || {
        let mut w = l2.write().unwrap();
        *w += 1;
    });
    thread::yield_now();

    drop(r1);
    assert_eq!(*r2, 0);
    thread::yield_now();
    thread::yield_now();
    thread::yield_now();
    assert_eq!(*r2, 0);
    drop(r2);
}

fn check_rwlock_unlock_bug2() {
    // There was a bug where when un-read-locking an rwlock by letting the last reader leaver,
    // we'd forget to wake up a writer.
    // That meant the writer thread could never run again.
    let l = Arc::new(RwLock::new(0));

    let r = l.read().unwrap();

    // Make a waiting writer.
    let l2 = l.clone();
    let h = thread::spawn(move || {
        let _w = l2.write().unwrap();
    });
    thread::yield_now();

    drop(r);
    h.join().unwrap();
}

fn park_timeout() {
    let start = Instant::now();

    thread::park_timeout(Duration::from_millis(200));
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn park_unpark() {
    let t1 = thread::current();
    let t2 = thread::spawn(move || {
        thread::park();
        thread::sleep(Duration::from_millis(200));
        t1.unpark();
    });

    let start = Instant::now();

    t2.thread().unpark();
    thread::park();
    // Normally, waiting in park/park_timeout may spuriously wake up early, but we
    // know Miri's timed synchronization primitives do not do that.

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn check_condvar() {
    let _ = std::sync::Condvar::new();
}

fn main() {
    check_barriers();
    check_conditional_variables_notify_one();
    check_conditional_variables_notify_all();
    check_conditional_variables_timed_wait_timeout();
    check_conditional_variables_timed_wait_notimeout();
    check_mutex();
    check_rwlock_write();
    check_rwlock_read_no_deadlock();
    check_once();
    check_rwlock_unlock_bug1();
    check_rwlock_unlock_bug2();
    park_timeout();
    park_unpark();
    check_condvar();
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::runtime::bootstrap;
use libra_config::utils;
use libra_mempool::mocks::MockSharedMempool;
use libra_proptest_helpers::ValueGenerator;
use libra_temppath::TempPath;
use libra_types::transaction::SignedTransaction;
use libradb::LibraDB;
use reqwest::blocking::Client;
use std::sync::Arc;

#[test]
fn test_json_rpc_service_fuzzer() {
    let mut gen = ValueGenerator::new();
    let data = generate_corpus(&mut gen);
    fuzzer(&data);
}

/// generate_corpus produces an arbitrary transaction to submit to JSON RPC service
pub fn generate_corpus(gen: &mut ValueGenerator) -> Vec<u8> {
    // use proptest to generate a SignedTransaction
    let txn = gen.generate(proptest::arbitrary::any::<SignedTransaction>());
    let payload = hex::encode(lcs::to_bytes(&txn).unwrap());
    let request =
        serde_json::json!({"jsonrpc": "2.0", "method": "submit", "params": [payload], "id": 1});
    serde_json::to_vec(&request).expect("failed to convert JSON to byte array")
}

pub fn fuzzer(data: &[u8]) {
    // set up mock Shared Mempool
    let smp = MockSharedMempool::new(None);

    // set up JSON RPC service
    let tmp_dir = TempPath::new();
    let db = LibraDB::new(&tmp_dir);
    let port = utils::get_available_port();
    let address = format!("0.0.0.0:{}", port);
    let _runtime = bootstrap(address.parse().unwrap(), Arc::new(db), smp.ac_client);

    let client = Client::new();
    let url = format!("http://{}", address);
    let json_request;
    match serde_json::from_slice::<serde_json::Value>(data) {
        Err(_) => return, // should not throw error or panic on invalid fuzzer inputs
        Ok(request) => json_request = request,
    };

    let response = client
        .post(&url) // address
        .json(&json_request)
        .send()
        .expect("failed to send request to JSON RPC server");

    let response: serde_json::Value = response.json().expect("failed to convert response to JSON");

    let json_rpc_protocol = response.get("jsonrpc");
    assert_eq!(
        json_rpc_protocol,
        Some(&serde_json::Value::String("2.0".to_string())),
        "JSON RPC response with incorrect protocol: {:?}",
        json_rpc_protocol
    );
    assert!(
        response.get("error").is_none(),
        "expected no error in JSON RPC response",
    );

    if let Some(result) = response.get("result") {
        let response_id: u64 = serde_json::from_value(response.get("id").unwrap().clone())
            .expect("failed to get ID from response");
        assert_eq!(response_id, 1, "mismatch ID in JSON RPC");
        assert!(
            *result == serde_json::Value::Null,
            "Received unexpected result payload from txn submission: {:?}",
            result
        );
    } else {
        panic!("Received JSON RPC response with no result payload");
    }
}
<commit_msg>[json-rpc] update fuzzing logic<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::runtime::bootstrap;
use libra_config::utils;
use libra_mempool::mocks::MockSharedMempool;
use libra_proptest_helpers::ValueGenerator;
use libra_temppath::TempPath;
use libra_types::transaction::SignedTransaction;
use libradb::LibraDB;
use reqwest::blocking::Client;
use std::sync::Arc;

#[test]
fn test_json_rpc_service_fuzzer() {
    let mut gen = ValueGenerator::new();
    let data = generate_corpus(&mut gen);
    fuzzer(&data);
}

/// generate_corpus produces an arbitrary transaction to submit to JSON RPC service
pub fn generate_corpus(gen: &mut ValueGenerator) -> Vec<u8> {
    // use proptest to generate a SignedTransaction
    let txn = gen.generate(proptest::arbitrary::any::<SignedTransaction>());
    let payload = hex::encode(lcs::to_bytes(&txn).unwrap());
    let request =
        serde_json::json!({"jsonrpc": "2.0", "method": "submit", "params": [payload], "id": 1});
    serde_json::to_vec(&request).expect("failed to convert JSON to byte array")
}

pub fn fuzzer(data: &[u8]) {
    // set up mock Shared Mempool
    let smp = MockSharedMempool::new(None);

    // set up JSON RPC service
    let tmp_dir = TempPath::new();
    let db = LibraDB::new(&tmp_dir);
    let port = utils::get_available_port();
    let address = format!("0.0.0.0:{}", port);
    let _runtime = bootstrap(address.parse().unwrap(), Arc::new(db), smp.ac_client);

    let client = Client::new();
    let url = format!("http://{}", address);
    let json_request = match serde_json::from_slice::<serde_json::Value>(data) {
        Err(_) => {
            // should not throw error or panic on invalid fuzzer inputs
            if cfg!(test) {
                panic!();
            }
            return;
        }
        Ok(request) => request,
    };

    let response = client
        .post(&url) // address
        .json(&json_request)
        .send()
        .expect("failed to send request to JSON RPC server");

    let response: serde_json::Value = response.json().expect("failed to convert response to JSON");

    let json_rpc_protocol = response.get("jsonrpc");
    assert_eq!(
        json_rpc_protocol,
        Some(&serde_json::Value::String("2.0".to_string())),
        "JSON RPC response with incorrect protocol: {:?}",
        json_rpc_protocol
    );
    if response.get("error").is_some() && cfg!(test) {
        panic!(
            "Unexpected error payload in JSON RPC response: {}",
            response
        );
    }

    // only proceed to check successful response for input from `generate_corpus`
    if !cfg!(test) {
        return;
    }

    let result = response.get("result").unwrap_or_else(|| {
        panic!(
            "Received JSON RPC response with no result payload. Full response: {}",
            response
        )
    });
    let response_id: u64 = serde_json::from_value(
        response
            .get("id")
            .unwrap_or_else(|| panic!("failed to get ID from response: {}", response))
            .clone(),
    )
    .unwrap_or_else(|_| panic!("Failed to deserialize ID from: {}", response));
    assert_eq!(response_id, 1, "mismatch ID in JSON RPC: {}", response);
    assert!(
        *result == serde_json::Value::Null,
        "Received unexpected result payload from txn submission: {:?}",
        response
    );
}
<|endoftext|>


<commit_before>//! The `optimizer` module contains iterators that allow a user to iterate over
//! `Program`s in order of increasing content
//!
//! The content is a measure of how complex a Program is.
use super::TIS_100::{Program, Instruction, Source, Destination, Register};

/// Allows one to determine the content of certain constructs, e.g. `Program`s
/// or `Instruction`s
pub trait Content {
    /// the content of the construct
    fn content(&self) -> u32;
}

impl Content for Program {
    fn content(&self) -> u32 {
        let Program(ref instructions) = *self;

        instructions.iter().fold(0 as u32, |sum, instruction| sum + instruction.content())
    }
}

impl Content for Instruction {
    fn content(&self) -> u32 {
        match *self {
            Instruction::NOP => 1,
            Instruction::MOV(ref source, ref destination) => 1 + source.content() + destination.content(),
            Instruction::SWP => 1,
            Instruction::SAV => 1,
            Instruction::ADD(ref source) => 1 + source.content(),
            Instruction::SUB(ref source) => 1 + source.content(),
        }
    }
}

impl Content for Source {
    fn content(&self) -> u32 {
        match *self {
            Source::Literal(value) => 1 + (value.abs() as u32),
            _ => 1,
        }
    }
}

impl Content for Destination {
    fn content(&self) -> u32 {
        1
    }
}

/// Iterator over Destination
pub struct DestinationIterator {
    current: i32,
    register_iterator: RegisterIterator,
}

impl DestinationIterator {
    /// create a DestinationIterator
    fn new() -> DestinationIterator {
        DestinationIterator {
            current: 0,
            register_iterator: RegisterIterator::new(),
        }
    }
}

impl Iterator for DestinationIterator {
    type Item = Destination;

    fn next(&mut self) -> Option<Destination> {
        let next_destination = match self.current {
            0 => Some(Destination::Port),
            _ => {
                self.register_iterator.next().map(|register| Destination::Register(register))
            }
        };
        self.current += 1;
        next_destination
    }
}

/// Iterator over Registers
pub struct RegisterIterator {
    current: u32,
}

impl RegisterIterator {
    fn new() -> RegisterIterator {
        RegisterIterator { current: 0 }
    }
}

impl Iterator for RegisterIterator {
    type Item = Register;

    fn next(&mut self) -> Option<Register> {
        let result = match self.current {
            0 => Some(Register::NIL),
            1 => Some(Register::ACC),
            _ => None
        };
        self.current += 1;
        result
    }
}

#[cfg(test)]
mod tests {
    use super::super::TIS_100::{Destination,Register};
    use super::*;

    #[test]
    fn should_iterator_over_destinations() {
        let destinations: Vec<Destination> = DestinationIterator::new().collect();

        assert_eq!(Destination::Port, destinations[0]);
        assert_eq!(Destination::Register(Register::NIL), destinations[1]);
        assert_eq!(Destination::Register(Register::ACC), destinations[2]);
    }

    #[test]
    fn should_iterator_over_registers() {
        let registers: Vec<Register> = RegisterIterator::new().collect();

        assert_eq!(Register::NIL, registers[0]);
        assert_eq!(Register::ACC, registers[1]);
    }
}
<commit_msg>Iterator over Programs<commit_after>//! The `optimizer` module contains iterators that allow a user to iterate over
//! `Program`s in order of increasing content
//!
//! The content is a measure of how complex a Program is.

use std::ops::{Rem, Div};
use super::TIS_100::{Program, Instruction, Source, Destination, Register};

/// Allows one to determine the content of certain constructs, e.g. `Program`s
/// or `Instruction`s
pub trait Content {
    /// the content of the construct
    fn content(&self) -> u32;
}

impl Content for Program {
    fn content(&self) -> u32 {
        let Program(ref instructions) = *self;

        instructions.iter().fold(0 as u32, |sum, instruction| sum + instruction.content())
    }
}

impl Content for Instruction {
    fn content(&self) -> u32 {
        match *self {
            Instruction::NOP => 1,
            Instruction::MOV(ref source, ref destination) => 1 + source.content() + destination.content(),
            Instruction::SWP => 1,
            Instruction::SAV => 1,
            Instruction::ADD(ref source) => 1 + source.content(),
            Instruction::SUB(ref source) => 1 + source.content(),
        }
    }
}

impl Content for Source {
    fn content(&self) -> u32 {
        match *self {
            Source::Literal(value) => 1 + (value.abs() as u32),
            _ => 1,
        }
    }
}

impl Content for Destination {
    fn content(&self) -> u32 {
        1
    }
}

/// Iterator over Programs
pub struct ProgramIterator {
    current: u32,
}

impl ProgramIterator {
    /// Create a `ProgramIterator`
    pub fn new() -> ProgramIterator {
        ProgramIterator { current: 0 }
    }
}

impl Iterator for ProgramIterator {
    type Item = Program;

    fn next(&mut self) -> Option<Program> {
        let instructions: Vec<Instruction> =
            digits_of(self.current, 33)
            .iter()
            .map(|digit| instruction_map(*digit))
            .collect();
        self.current += 1;
        Some(Program(instructions))
    }
}

fn digits_of(mut n: u32, base: u32) -> Vec<u32> {
    let mut digits: Vec<u32> = vec![];
    while n > 0 {
        digits.push(n.rem(base));
        n = n.div(base);
    }
    digits
}

fn instruction_map(index: u32) -> Instruction {
    match index % 33 {
        0 => Instruction::NOP,
        1 => Instruction::SWP,
        2 => Instruction::SAV,
        3 => Instruction::ADD(Source::Port),
        4 => Instruction::ADD(Source::Register(Register::NIL)),
        5 => Instruction::ADD(Source::Register(Register::ACC)),
        6 => Instruction::ADD(Source::Literal(0)),
        7 => Instruction::ADD(Source::Literal(1)),
        8 => Instruction::ADD(Source::Literal(-1)),
        9 => Instruction::SUB(Source::Port),
        10 => Instruction::SUB(Source::Register(Register::NIL)),
        11 => Instruction::SUB(Source::Register(Register::ACC)),
        12 => Instruction::SUB(Source::Literal(0)),
        13 => Instruction::SUB(Source::Literal(1)),
        14 => Instruction::SUB(Source::Literal(-1)),
        15 => Instruction::MOV(Source::Port, Destination::Port),
        16 => Instruction::MOV(Source::Register(Register::NIL), Destination::Port),
        17 => Instruction::MOV(Source::Register(Register::ACC), Destination::Port),
        18 => Instruction::MOV(Source::Literal(0), Destination::Port),
        19 => Instruction::MOV(Source::Literal(1), Destination::Port),
        20 => Instruction::MOV(Source::Literal(-1), Destination::Port),
        21 => Instruction::MOV(Source::Port, Destination::Register(Register::NIL)),
        22 => Instruction::MOV(Source::Register(Register::NIL), Destination::Register(Register::NIL)),
        23 => Instruction::MOV(Source::Register(Register::ACC), Destination::Register(Register::NIL)),
        24 => Instruction::MOV(Source::Literal(0), Destination::Register(Register::NIL)),
        25 => Instruction::MOV(Source::Literal(1), Destination::Register(Register::NIL)),
        26 => Instruction::MOV(Source::Literal(-1), Destination::Register(Register::NIL)),
        27 => Instruction::MOV(Source::Port, Destination::Register(Register::ACC)),
        28 => Instruction::MOV(Source::Register(Register::NIL), Destination::Register(Register::ACC)),
        29 => Instruction::MOV(Source::Register(Register::ACC), Destination::Register(Register::ACC)),
        30 => Instruction::MOV(Source::Literal(0), Destination::Register(Register::ACC)),
        31 => Instruction::MOV(Source::Literal(1), Destination::Register(Register::ACC)),
        32 => Instruction::MOV(Source::Literal(-1), Destination::Register(Register::ACC)),
        _ => Instruction::NOP,
    }
}

/// Iterator over Destination
pub struct DestinationIterator {
    current: i32,
    register_iterator: RegisterIterator,
}

impl DestinationIterator {
    /// create a DestinationIterator
    fn new() -> DestinationIterator {
        DestinationIterator {
            current: 0,
            register_iterator: RegisterIterator::new(),
        }
    }
}

impl Iterator for DestinationIterator {
    type Item = Destination;

    fn next(&mut self) -> Option<Destination> {
        let next_destination = match self.current {
            0 => Some(Destination::Port),
            _ => {
                self.register_iterator.next().map(|register| Destination::Register(register))
            }
        };
        self.current += 1;
        next_destination
    }
}

/// Iterator over Registers
pub struct RegisterIterator {
    current: u32,
}

impl RegisterIterator {
    fn new() -> RegisterIterator {
        RegisterIterator { current: 0 }
    }
}

impl Iterator for RegisterIterator {
    type Item = Register;

    fn next(&mut self) -> Option<Register> {
        let result = match self.current {
            0 => Some(Register::NIL),
            1 => Some(Register::ACC),
            _ => None
        };
        self.current += 1;
        result
    }
}

#[cfg(test)]
mod tests {
    use super::super::TIS_100::{Program,Instruction,Source,Destination,Register};
    use super::*;

    #[test]
    fn should_iterator_over_destinations() {
        let destinations: Vec<Destination> = DestinationIterator::new().collect();

        assert_eq!(Destination::Port, destinations[0]);
        assert_eq!(Destination::Register(Register::NIL), destinations[1]);
        assert_eq!(Destination::Register(Register::ACC), destinations[2]);
    }

    #[test]
    fn should_iterator_over_registers() {
        let registers: Vec<Register> = RegisterIterator::new().collect();

        assert_eq!(Register::NIL, registers[0]);
        assert_eq!(Register::ACC, registers[1]);
    }

    #[test]
    fn should_iterate_over_programs() {
        let programs: Vec<Program> = ProgramIterator::new().take(34).collect();

        assert_eq!(Program(vec![]), programs[0]);
        assert_eq!(Program(vec![Instruction::SWP]), programs[1]);
        assert_eq!(Program(vec![Instruction::SAV]), programs[2]);
        assert_eq!(Program(vec![
            Instruction::MOV(
                Source::Literal(-1),
                Destination::Register(Register::ACC)),
        ]), programs[32]);
        assert_eq!(Program(vec![
            Instruction::NOP,
            Instruction::SWP,
        ]), programs[33])
    }
}
<|endoftext|>