<commit_before>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;
use url::Url;

use Provider;
use Track;
use Playlist;
use open_graph;
use apple_music;
use youtube;
use soundcloud;
use spotify;
use error::Error;

use url::percent_encoding::{percent_decode};

static APPLE_MUSIC_SONG:     &'static str = r"tools.applemusic.com/embed/v1/song/([a-zA-Z0-9_-]+)";
static APPLE_MUSIC_ALBUM:    &'static str = r"tools.applemusic.com/embed/v1/album/([a-zA-Z0-9_-]+)";
static APPLE_MUSIC_PLAYLIST: &'static str = r"tools.applemusic.com/embed/v1/playlist/pl.([a-zA-Z0-9_-]+)";
static YOUTUBE_EMBED:        &'static str = r"www.youtube.com/embed/([a-zA-Z0-9_-].+)";
static YOUTUBE_LIST:         &'static str = r"www.youtube.com/embed/videoseries\?list=([a-zA-Z0-9_-]+)";
static YOUTUBE_WATCH:        &'static str = r"www.youtube.com/watch\?v=([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_TRACK:     &'static str = r"api.soundcloud.com/tracks/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_PLAYLIST:  &'static str = r"api.soundcloud.com/playlists/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_USER:      &'static str = r"api.soundcloud.com/users/([a-zA-Z0-9_-]+)";
static SPOTIFY_TRACK:        &'static str = r"open.spotify.com/track/([a-zA-Z0-9_-]+)";
static SPOTIFY_TRACK_OPEN:   &'static str = r"spotify:track:([a-zA-Z0-9_-]+)";
static SPOTIFY_PLAYLIST:     &'static str = r"(spotify:user:[a-zA-Z0-9_-]+:playlist:[a-zA-Z0-9_-]+)";

#[derive(Debug)]
pub struct ScraperProduct {
    pub playlists: Vec<Playlist>,
    pub tracks:    Vec<Track>,
    pub og_obj:    Option<open_graph::Object>,
}

pub fn extract(url: &str) -> Result<ScraperProduct, Error> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut playlists = Vec::new();
        let mut tracks    = Vec::new();
        let mut og_props  = Vec::new();
        walk(dom.document, &mut playlists, &mut tracks, &mut og_props);
        let og_obj = if og_props.len() > 0 {
            Some(open_graph::Object::new(&og_props))
        } else {
            None
        };
        Ok(ScraperProduct {
            playlists: playlists,
            tracks:    tracks,
            og_obj:    og_obj
        })
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        Err(Error::NotFound)
    }
}

// This is not proper HTML serialization, of course.
fn walk(handle:    Handle,
        playlists: &mut Vec<Playlist>,
        tracks:    &mut Vec<Track>,
        og_props:  &mut Vec<(String, String)>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let mut ps = extract_open_graph_metadata_from_tag(tag_name, attrs);
            og_props.append(&mut ps);
            let (ps, ts) = extract_enclosures_from_tag(tag_name, attrs);
            for playlist in ps.iter().cloned() {
                if !(playlists).iter().any(|p| playlist == *p) {
                    (*playlists).push(playlist)
                }
            }
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(child.clone(), playlists, tracks, og_props);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

pub fn extract_enclosures_from_tag(tag_name: &str,
                                   attrs: &Vec<Attribute>) -> (Vec<Playlist>, Vec<Track>) {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_enclosures_from_url(src.to_string()),
            None => (vec![], vec![])
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_enclosures_from_url(href.to_string()),
            None => (vec![], vec![])
        }
    } else {
        (vec![], vec![])
    }
}

pub fn extract_open_graph_metadata_from_tag(tag_name: &str,
                                            attrs: &Vec<Attribute>) -> Vec<(String, String)> {

    let mut og_props = vec!();
    if tag_name == "meta" {
        match extract_open_graph_prop("property", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
        match extract_open_graph_prop("name", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
    }
    og_props
}

fn extract_open_graph_prop<'a>(attr_name: &str, attrs: &Vec<Attribute>) -> Option<(String, String)> {
    attr(attr_name, attrs)
        .and_then(|property|
                  if property.starts_with("og:") {
                      let end = property.chars().count();
                      let key = unsafe {
                          property.slice_unchecked(3, end)
                      }.to_string();
                      attr("content", attrs).map(|content| (key, content))
                  } else {
                      None
                  })
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => {
                let strs: Vec<&str> = cap[1].split('?').collect();
                return Some(strs[0].to_string())
            },
            None => None
        },
        Err(_) => None
    }
}

fn country_param(url: &str) -> String {
    let url = Url::parse(&url).unwrap();
    if let Some(q) = url.query() {
        let strs: Vec<&str> = q.split('=').collect();
        if strs[0] == "country" {
            return strs[1].to_string();
        }
    };
    "en".to_string()
}

fn fetch_spotify_playlist(uri: &str) -> (Vec<Playlist>, Vec<Track>) {
    match Regex::new(r"spotify:user:([a-zA-Z0-9_-]+):playlist:([a-zA-Z0-9_-]+)") {
        Ok(re) => match re.captures(uri) {
            Some(cap) => {
                let user_id     = cap[1].to_string();
                let playlist_id = cap[2].to_string();
                return match spotify::fetch_playlist(&user_id, &playlist_id) {
                    Ok(playlist) => {
                        let tracks = playlist.tracks.items
                            .iter()
                            .map(|ref i| Track::from_sp_track(&i.track))
                            .collect::<Vec<_>>();
                        return (vec![Playlist::from_sp_playlist(&playlist)], tracks);
                    },
                    Err(_)       => (vec![], vec![])
                }
            },
            None => (vec![], vec![])
        },
        Err(_) => (vec![], vec![])
    }
}

fn fetch_spotify_track(identifier: String) -> Vec<Track> {
    match spotify::fetch_track(&identifier) {
        Ok(t) => {
            let ref mut track = Track::new(Provider::Spotify, identifier);
            track.update_with_sp_track(&t);
            vec![track.clone()]
        },
        Err(_) => vec![Track::new(Provider::Spotify, identifier)],
    }
}

fn fetch_youtube_playlist(id: &str) -> (Vec<Playlist>, Vec<Track>) {
    let playlists = match youtube::fetch_playlist(id) {
        Ok(res) => res.items
            .iter()
            .map(|ref i| Playlist::from_yt_playlist(i))
            .collect::<Vec<_>>(),
        Err(_)  => vec![],
    };
    let tracks = match youtube::fetch_playlist_items(id) {
        Ok(res) => res.items
            .iter()
            .map(|ref i| Track::from_yt_playlist_item(i))
            .collect::<Vec<_>>(),
        Err(_) => vec![],
    };
    (playlists, tracks)
}

fn extract_enclosures_from_url(url: String) -> (Vec<Playlist>, Vec<Track>) {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();

    match extract_identifier(&decoded, APPLE_MUSIC_SONG) {
        Some(identifier) => {
            let country = country_param(&url);
            if let Ok(song) = apple_music::fetch_song(&identifier, &country) {
                return (vec![], vec![Track::from_am_song(&song)])
            };
        },
        None => ()
    }
    match extract_identifier(&decoded, APPLE_MUSIC_ALBUM) {
        Some(identifier) => {
            // todo
            return (vec![], vec![])
        },
        None => ()
    }
    match extract_identifier(&decoded, APPLE_MUSIC_PLAYLIST) {
        Some(identifier) => {
            let country = country_param(&url);
            if let Ok(playlist) = apple_music::fetch_playlist(&identifier, &country) {
                return (vec![Playlist::from_am_playlist(&playlist)], vec![])
            }
            return (vec![], vec![])
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => return (vec![], vec![Track::new(Provider::YouTube, identifier)]),
        None             => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => return fetch_youtube_playlist(&identifier),
        None             => ()
    }
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => return (vec![], vec![Track::new(Provider::YouTube, identifier)]),
        None             => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => return (vec![], vec![Track::new(Provider::SoundCloud, identifier)]),
        None             => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => return match soundcloud::fetch_playlist(&identifier) {
            Ok(playlist) => {
                let tracks = playlist.tracks
                    .iter()
                    .map(|ref t| Track::from_sc_track(t))
                    .collect::<Vec<_>>();
                (vec![Playlist::from_sc_playlist(&playlist)], tracks)
            },
            Err(_)       => (vec![], vec![]),
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => return match soundcloud::fetch_user_tracks(&identifier) {
            Ok(tracks) => {
                let tracks = tracks
                    .iter()
                    .map(|ref t| Track::from_sc_track(t))
                    .collect::<Vec<_>>();
                (vec![], tracks)
            },
            Err(_)       => (vec![], vec![]),
        },
        None => ()
    }
    match extract_identifier(&decoded, SPOTIFY_TRACK) {
        Some(identifier) => return (vec![], fetch_spotify_track(identifier)),
        None             => ()
    }
    match extract_identifier(&decoded, SPOTIFY_TRACK_OPEN) {
        Some(identifier) => return (vec![], fetch_spotify_track(identifier)),
        None             => ()
    }
    match extract_identifier(&decoded, SPOTIFY_PLAYLIST) {
        Some(uri) => return fetch_spotify_playlist(&uri),
        None      => ()
    }
    return (vec![], vec![])
}

#[cfg(test)]
mod test {
    use super::extract;
    use super::extract_identifier;
    use Provider;
    use Track;

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494/stream&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }
    }
    #[test]
    fn test_extract() {
        let url       = "http://spincoaster.com/spincoaster-breakout-2017";
        let product   = extract(url).unwrap();
        let playlists = product.playlists;
        let tracks    = product.tracks;
        assert_eq!(playlists.len(), 2);
        let youtube_tracks: Vec<&Track> = tracks.iter().filter(|&x| x.provider == Provider::YouTube).collect();
        let spotify_tracks: Vec<&Track> = tracks.iter().filter(|&x| x.provider == Provider::Spotify).collect();
        assert_eq!(youtube_tracks.len(), 15);
        assert_eq!(spotify_tracks.len(), 30);
    }
}
<commit_msg>Fix scraper::country_param: Check if url has protocol<commit_after>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;
use url::Url;

use Provider;
use Track;
use Playlist;
use open_graph;
use apple_music;
use youtube;
use soundcloud;
use spotify;
use error::Error;

use url::percent_encoding::{percent_decode};

static APPLE_MUSIC_SONG:     &'static str = r"tools.applemusic.com/embed/v1/song/([a-zA-Z0-9_-]+)";
static APPLE_MUSIC_ALBUM:    &'static str = r"tools.applemusic.com/embed/v1/album/([a-zA-Z0-9_-]+)";
static APPLE_MUSIC_PLAYLIST: &'static str = r"tools.applemusic.com/embed/v1/playlist/pl.([a-zA-Z0-9_-]+)";
static YOUTUBE_EMBED:        &'static str = r"www.youtube.com/embed/([a-zA-Z0-9_-].+)";
static YOUTUBE_LIST:         &'static str = r"www.youtube.com/embed/videoseries\?list=([a-zA-Z0-9_-]+)";
static YOUTUBE_WATCH:        &'static str = r"www.youtube.com/watch\?v=([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_TRACK:     &'static str = r"api.soundcloud.com/tracks/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_PLAYLIST:  &'static str = r"api.soundcloud.com/playlists/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_USER:      &'static str = r"api.soundcloud.com/users/([a-zA-Z0-9_-]+)";
static SPOTIFY_TRACK:        &'static str = r"open.spotify.com/track/([a-zA-Z0-9_-]+)";
static SPOTIFY_TRACK_OPEN:   &'static str = r"spotify:track:([a-zA-Z0-9_-]+)";
static SPOTIFY_PLAYLIST:     &'static str = r"(spotify:user:[a-zA-Z0-9_-]+:playlist:[a-zA-Z0-9_-]+)";

#[derive(Debug)]
pub struct ScraperProduct {
    pub playlists: Vec<Playlist>,
    pub tracks:    Vec<Track>,
    pub og_obj:    Option<open_graph::Object>,
}

pub fn extract(url: &str) -> Result<ScraperProduct, Error> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut playlists = Vec::new();
        let mut tracks    = Vec::new();
        let mut og_props  = Vec::new();
        walk(dom.document, &mut playlists, &mut tracks, &mut og_props);
        let og_obj = if og_props.len() > 0 {
            Some(open_graph::Object::new(&og_props))
        } else {
            None
        };
        Ok(ScraperProduct {
            playlists: playlists,
            tracks:    tracks,
            og_obj:    og_obj
        })
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        Err(Error::NotFound)
    }
}

// This is not proper HTML serialization, of course.
fn walk(handle:    Handle,
        playlists: &mut Vec<Playlist>,
        tracks:    &mut Vec<Track>,
        og_props:  &mut Vec<(String, String)>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let mut ps = extract_open_graph_metadata_from_tag(tag_name, attrs);
            og_props.append(&mut ps);
            let (ps, ts) = extract_enclosures_from_tag(tag_name, attrs);
            for playlist in ps.iter().cloned() {
                if !(playlists).iter().any(|p| playlist == *p) {
                    (*playlists).push(playlist)
                }
            }
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(child.clone(), playlists, tracks, og_props);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

pub fn extract_enclosures_from_tag(tag_name: &str,
                                   attrs: &Vec<Attribute>) -> (Vec<Playlist>, Vec<Track>) {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_enclosures_from_url(src.to_string()),
            None => (vec![], vec![])
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_enclosures_from_url(href.to_string()),
            None => (vec![], vec![])
        }
    } else {
        (vec![], vec![])
    }
}

pub fn extract_open_graph_metadata_from_tag(tag_name: &str,
                                            attrs: &Vec<Attribute>) -> Vec<(String, String)> {

    let mut og_props = vec!();
    if tag_name == "meta" {
        match extract_open_graph_prop("property", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
        match extract_open_graph_prop("name", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
    }
    og_props
}

fn extract_open_graph_prop<'a>(attr_name: &str, attrs: &Vec<Attribute>) -> Option<(String, String)> {
    attr(attr_name, attrs)
        .and_then(|property|
                  if property.starts_with("og:") {
                      let end = property.chars().count();
                      let key = unsafe {
                          property.slice_unchecked(3, end)
                      }.to_string();
                      attr("content", attrs).map(|content| (key, content))
                  } else {
                      None
                  })
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => {
                let strs: Vec<&str> = cap[1].split('?').collect();
                return Some(strs[0].to_string())
            },
            None => None
        },
        Err(_) => None
    }
}

fn country_param(url_str: &str) -> String {
    let url_str = if url_str.starts_with("http") || url_str.starts_with("https") {
        url_str.to_string()
    } else {
        "https://".to_string() + url_str
    };
    let url = Url::parse(&url_str).unwrap();
    if let Some(q) = url.query() {
        let strs: Vec<&str> = q.split('=').collect();
        if strs[0] == "country" {
            return strs[1].to_string();
        }
    };
    "us".to_string()
}

fn fetch_spotify_playlist(uri: &str) -> (Vec<Playlist>, Vec<Track>) {
    match Regex::new(r"spotify:user:([a-zA-Z0-9_-]+):playlist:([a-zA-Z0-9_-]+)") {
        Ok(re) => match re.captures(uri) {
            Some(cap) => {
                let user_id     = cap[1].to_string();
                let playlist_id = cap[2].to_string();
                return match spotify::fetch_playlist(&user_id, &playlist_id) {
                    Ok(playlist) => {
                        let tracks = playlist.tracks.items
                            .iter()
                            .map(|ref i| Track::from_sp_track(&i.track))
                            .collect::<Vec<_>>();
                        return (vec![Playlist::from_sp_playlist(&playlist)], tracks);
                    },
                    Err(_)       => (vec![], vec![])
                }
            },
            None => (vec![], vec![])
        },
        Err(_) => (vec![], vec![])
    }
}

fn fetch_spotify_track(identifier: String) -> Vec<Track> {
    match spotify::fetch_track(&identifier) {
        Ok(t) => {
            let ref mut track = Track::new(Provider::Spotify, identifier);
            track.update_with_sp_track(&t);
            vec![track.clone()]
        },
        Err(_) => vec![Track::new(Provider::Spotify, identifier)],
    }
}

fn fetch_youtube_playlist(id: &str) -> (Vec<Playlist>, Vec<Track>) {
    let playlists = match youtube::fetch_playlist(id) {
        Ok(res) => res.items
            .iter()
            .map(|ref i| Playlist::from_yt_playlist(i))
            .collect::<Vec<_>>(),
        Err(_)  => vec![],
    };
    let tracks = match youtube::fetch_playlist_items(id) {
        Ok(res) => res.items
            .iter()
            .map(|ref i| Track::from_yt_playlist_item(i))
            .collect::<Vec<_>>(),
        Err(_) => vec![],
    };
    (playlists, tracks)
}

fn extract_enclosures_from_url(url: String) -> (Vec<Playlist>, Vec<Track>) {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();

    match extract_identifier(&decoded, APPLE_MUSIC_SONG) {
        Some(identifier) => {
            let country = country_param(&url);
            if let Ok(song) = apple_music::fetch_song(&identifier, &country) {
                return (vec![], vec![Track::from_am_song(&song)])
            };
        },
        None => ()
    }
    match extract_identifier(&decoded, APPLE_MUSIC_ALBUM) {
        Some(identifier) => {
            // todo
            return (vec![], vec![])
        },
        None => ()
    }
    match extract_identifier(&decoded, APPLE_MUSIC_PLAYLIST) {
        Some(identifier) => {
            let country = country_param(&url);
            match apple_music::fetch_playlist(&identifier, &country) {
                Ok(playlist) => return (vec![Playlist::from_am_playlist(&playlist)], vec![]),
                Err(_) => (),
            }
            return (vec![], vec![])
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => return (vec![], vec![Track::new(Provider::YouTube, identifier)]),
        None             => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => return fetch_youtube_playlist(&identifier),
        None             => ()
    }
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => return (vec![], vec![Track::new(Provider::YouTube, identifier)]),
        None             => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => return (vec![], vec![Track::new(Provider::SoundCloud, identifier)]),
        None             => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => return match soundcloud::fetch_playlist(&identifier) {
            Ok(playlist) => {
                let tracks = playlist.tracks
                    .iter()
                    .map(|ref t| Track::from_sc_track(t))
                    .collect::<Vec<_>>();
                (vec![Playlist::from_sc_playlist(&playlist)], tracks)
            },
            Err(_)       => (vec![], vec![]),
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => return match soundcloud::fetch_user_tracks(&identifier) {
            Ok(tracks) => {
                let tracks = tracks
                    .iter()
                    .map(|ref t| Track::from_sc_track(t))
                    .collect::<Vec<_>>();
                (vec![], tracks)
            },
            Err(_)       => (vec![], vec![]),
        },
        None => ()
    }
    match extract_identifier(&decoded, SPOTIFY_TRACK) {
        Some(identifier) => return (vec![], fetch_spotify_track(identifier)),
        None             => ()
    }
    match extract_identifier(&decoded, SPOTIFY_TRACK_OPEN) {
        Some(identifier) => return (vec![], fetch_spotify_track(identifier)),
        None             => ()
    }
    match extract_identifier(&decoded, SPOTIFY_PLAYLIST) {
        Some(uri) => return fetch_spotify_playlist(&uri),
        None      => ()
    }
    return (vec![], vec![])
}

#[cfg(test)]
mod test {
    use super::extract;
    use super::extract_identifier;
    use Provider;
    use Track;

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494/stream&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }
    }
    #[test]
    fn test_extract() {
        let url       = "http://spincoaster.com/spincoaster-breakout-2017";
        let product   = extract(url).unwrap();
        let playlists = product.playlists;
        let tracks    = product.tracks;
        assert_eq!(playlists.len(), 2);
        let youtube_tracks: Vec<&Track> = tracks.iter().filter(|&x| x.provider == Provider::YouTube).collect();
        let spotify_tracks: Vec<&Track> = tracks.iter().filter(|&x| x.provider == Provider::Spotify).collect();
        assert_eq!(youtube_tracks.len(), 15);
        assert_eq!(spotify_tracks.len(), 30);
    }
}
<|endoftext|>


<commit_before>use crate::utils;
use dirs;
use std::fs;
use std::io::ErrorKind;
use std::process::Command;

pub fn setup() {
    // Environment (Unix)

    utils::install_app("1password");
    utils::install_package("ripgrep");

    // Version control (Git)

    utils::install_package("git");

    let home_path = match dirs::home_dir() {
        Some(directory) => directory,
        None => panic!("Cannot find the home directory."),
    };

    let repos_path = home_path.join("repos");

    utils::create_dir("repos", &repos_path);

    match fs::read_dir(&repos_path) {
        Ok(_) => println!("The repos are already installed."),
        Err(_) => {
            println!("Cloning scripts repo.");

            match Command::new("git")
                .current_dir(&repos_path)
                .arg("clone")
                .arg("https://github.com/trevordmiller/scripts")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Cloning study repo.");

            match Command::new("git")
                .current_dir(&repos_path)
                .arg("clone")
                .arg("https://github.com/trevordmiller/study")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Cloning trevordmiller.github.io repo.");

            match Command::new("git")
                .current_dir(&repos_path)
                .arg("clone")
                .arg("https://github.com/trevordmiller/trevordmiller.github.io")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }
        }
    };

    // Editor (Vim)

    utils::install_package("vim");

    let editor_plugins_path = home_path
        .join(".vim")
        .join("pack")
        .join("plugins")
        .join("start");

    utils::create_dir("editor plugins", &editor_plugins_path);

    let editor_configuration_path = home_path.join(".vimrc");

    utils::create_file("editor configuration", &editor_configuration_path);

    match fs::read_dir(&editor_plugins_path) {
        Ok(_) => println!("The editor plugins are already installed."),
        Err(_) => {
            println!("Installing editor plugin to normalize editor defaults.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/tpope/vim-sensible")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor indentation.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/tpope/vim-sleuth")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor languages.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/sheerun/vim-polyglot")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor ignore patterns.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/octref/RootIgnore")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor static analysis.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/dense-analysis/ale")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Adding editor configuration to enhance editor search.");

            match fs::write(&editor_configuration_path, "set grepprg=rg\\ --vimgrep") {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            match fs::write(&editor_configuration_path, "set grepformat=%f:%l:%c:%m") {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor colors.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/arcticicestudio/nord-vim")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Adding editor configuration to enhance editor colors.");

            match fs::write(editor_configuration_path, "colorscheme nord") {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }
        }
    };

    // Application programming (JavaScript)

    utils::install_package("node");

    // Systems programming (Rust)

    utils::install_package("rustup-init");

    let rustup_path_check = match Command::new("which").arg("rustup").output() {
        Ok(output) => output.stdout,
        Err(error) => panic!("There was a problem: {:?}", error),
    };

    if rustup_path_check.is_empty() {
        match Command::new("rustup-init").output() {
            Ok(_) => (),
            Err(error) => match error.kind() {
                ErrorKind::NotFound => panic!("The rustup-init command is missing."),
                other_error => panic!("There was a problem: {:?}", other_error),
            },
        }
    } else {
        println!("The rustup-init has already been run.")
    }
}

pub fn upgrade() {
    // Environment (Unix)

    println!("Upgrading package manager.");

    match Command::new("brew").arg("update").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    println!("Upgrading password manager.");

    match Command::new("brew")
        .arg("cask")
        .arg("upgrade")
        .arg("1password")
        .output()
    {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    println!("Upgrading search tool.");

    match Command::new("brew").arg("upgrade").arg("ripgrep").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Version control (Git)

    println!("Upgrading version control.");

    match Command::new("brew").arg("upgrade").arg("git").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Editor (Vim)

    println!("Upgrading editor.");

    match Command::new("brew").arg("upgrade").arg("vim").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    println!("Upgrading editor plugins.");

    let home_path = match dirs::home_dir() {
        Some(directory) => directory,
        None => panic!("Cannot find the home directory."),
    };

    let editor_plugins_path = home_path
        .join(".vim")
        .join("pack")
        .join("plugins")
        .join("start");

    match fs::read_dir(&editor_plugins_path) {
        Ok(paths) => {
            for path in paths {
                match Command::new("git")
                    .current_dir(match path {
                        Ok(path) => path.path(),
                        Err(error) => panic!("There was a problem: {:?}", error),
                    })
                    .arg("pull")
                    .output()
                {
                    Ok(_) => (),
                    Err(error) => panic!("There was a problem: {:?}", error),
                }
            }
        }
        Err(error) => panic!("There was a problem: {:?}", error),
    };

    // Application programming (JavaScript)

    println!("Upgrading application programming toolchain.");

    match Command::new("brew").arg("upgrade").arg("node").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Systems programming (Rust)

    println!("Upgrading systems programming toolchain.");

    match Command::new("brew")
        .arg("upgrade")
        .arg("rustup-init")
        .output()
    {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match Command::new("rustup").arg("update").output() {
        Ok(_) => (),
        Err(error) => match error.kind() {
            ErrorKind::NotFound => panic!("The rustup command is missing."),
            other_error => panic!("There was a problem: {:?}", other_error),
        },
    }
}

pub fn end() {
    // Environment (Unix)

    println!("Removing package manager artifacts.");

    match Command::new("brew").arg("cleanup").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Version control (Git)

    println!("Making sure all repos are clean.");

    let home_path = match dirs::home_dir() {
        Some(directory) => directory,
        None => panic!("Cannot find the home directory."),
    };

    let repos_path = home_path.join("repos");

    match fs::read_dir(&repos_path) {
        Ok(paths) => {
            let mut all_repos_clean = true;

            for path in paths {
                let path = match &path {
                    Ok(path) => path.path(),
                    Err(error) => panic!("There was a problem: {:?}", error),
                };

                let status_check = match Command::new("git")
                    .current_dir(&path)
                    .arg("status")
                    .arg("--porcelain")
                    .output()
                {
                    Ok(output) => output.stdout,
                    Err(error) => panic!("There was a problem: {:?}", error),
                };

                let unpushed_check = match Command::new("git")
                    .current_dir(&path)
                    .arg("log")
                    .arg("@{u}..")
                    .output()
                {
                    Ok(output) => output.stdout,
                    Err(error) => panic!("There was a problem: {:?}", error),
                };

                if !status_check.is_empty() || !unpushed_check.is_empty() {
                    all_repos_clean = false;
                    eprintln!("A dirty repo was found: {}", &path.display());
                }
            }

            if !all_repos_clean {
                panic!("Repos are dirty.");
            }
        }
        Err(error) => panic!("There was a problem: {:?}", error),
    };

    // Editor (Vim)

    println!("Quitting editor processes.");

    match Command::new("killall").arg("vim").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Application programming (JavaScript)

    println!("Quitting application programming processes.");

    match Command::new("killall").arg("node").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Systems programming (Rust)

    println!("Quitting systems programming processes.");

    match Command::new("killall").arg("rls").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}
<commit_msg>Remove github pages website hosting from cloning<commit_after>use crate::utils;
use dirs;
use std::fs;
use std::io::ErrorKind;
use std::process::Command;

pub fn setup() {
    // Environment (Unix)

    utils::install_app("1password");
    utils::install_package("ripgrep");

    // Version control (Git)

    utils::install_package("git");

    let home_path = match dirs::home_dir() {
        Some(directory) => directory,
        None => panic!("Cannot find the home directory."),
    };

    let repos_path = home_path.join("repos");

    utils::create_dir("repos", &repos_path);

    match fs::read_dir(&repos_path) {
        Ok(_) => println!("The repos are already installed."),
        Err(_) => {
            println!("Cloning scripts repo.");

            match Command::new("git")
                .current_dir(&repos_path)
                .arg("clone")
                .arg("https://github.com/trevordmiller/scripts")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Cloning study repo.");

            match Command::new("git")
                .current_dir(&repos_path)
                .arg("clone")
                .arg("https://github.com/trevordmiller/study")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }
        }
    };

    // Editor (Vim)

    utils::install_package("vim");

    let editor_plugins_path = home_path
        .join(".vim")
        .join("pack")
        .join("plugins")
        .join("start");

    utils::create_dir("editor plugins", &editor_plugins_path);

    let editor_configuration_path = home_path.join(".vimrc");

    utils::create_file("editor configuration", &editor_configuration_path);

    match fs::read_dir(&editor_plugins_path) {
        Ok(_) => println!("The editor plugins are already installed."),
        Err(_) => {
            println!("Installing editor plugin to normalize editor defaults.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/tpope/vim-sensible")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor indentation.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/tpope/vim-sleuth")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor languages.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/sheerun/vim-polyglot")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor ignore patterns.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/octref/RootIgnore")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor static analysis.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/dense-analysis/ale")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Adding editor configuration to enhance editor search.");

            match fs::write(&editor_configuration_path, "set grepprg=rg\\ --vimgrep") {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            match fs::write(&editor_configuration_path, "set grepformat=%f:%l:%c:%m") {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Installing editor plugin to enhance editor colors.");

            match Command::new("git")
                .current_dir(&editor_plugins_path)
                .arg("clone")
                .arg("https://github.com/arcticicestudio/nord-vim")
                .output()
            {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            println!("Adding editor configuration to enhance editor colors.");

            match fs::write(editor_configuration_path, "colorscheme nord") {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }
        }
    };

    // Application programming (JavaScript)

    utils::install_package("node");

    // Systems programming (Rust)

    utils::install_package("rustup-init");

    let rustup_path_check = match Command::new("which").arg("rustup").output() {
        Ok(output) => output.stdout,
        Err(error) => panic!("There was a problem: {:?}", error),
    };

    if rustup_path_check.is_empty() {
        match Command::new("rustup-init").output() {
            Ok(_) => (),
            Err(error) => match error.kind() {
                ErrorKind::NotFound => panic!("The rustup-init command is missing."),
                other_error => panic!("There was a problem: {:?}", other_error),
            },
        }
    } else {
        println!("The rustup-init has already been run.")
    }
}

pub fn upgrade() {
    // Environment (Unix)

    println!("Upgrading package manager.");

    match Command::new("brew").arg("update").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    println!("Upgrading password manager.");

    match Command::new("brew")
        .arg("cask")
        .arg("upgrade")
        .arg("1password")
        .output()
    {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    println!("Upgrading search tool.");

    match Command::new("brew").arg("upgrade").arg("ripgrep").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Version control (Git)

    println!("Upgrading version control.");

    match Command::new("brew").arg("upgrade").arg("git").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Editor (Vim)

    println!("Upgrading editor.");

    match Command::new("brew").arg("upgrade").arg("vim").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    println!("Upgrading editor plugins.");

    let home_path = match dirs::home_dir() {
        Some(directory) => directory,
        None => panic!("Cannot find the home directory."),
    };

    let editor_plugins_path = home_path
        .join(".vim")
        .join("pack")
        .join("plugins")
        .join("start");

    match fs::read_dir(&editor_plugins_path) {
        Ok(paths) => {
            for path in paths {
                match Command::new("git")
                    .current_dir(match path {
                        Ok(path) => path.path(),
                        Err(error) => panic!("There was a problem: {:?}", error),
                    })
                    .arg("pull")
                    .output()
                {
                    Ok(_) => (),
                    Err(error) => panic!("There was a problem: {:?}", error),
                }
            }
        }
        Err(error) => panic!("There was a problem: {:?}", error),
    };

    // Application programming (JavaScript)

    println!("Upgrading application programming toolchain.");

    match Command::new("brew").arg("upgrade").arg("node").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Systems programming (Rust)

    println!("Upgrading systems programming toolchain.");

    match Command::new("brew")
        .arg("upgrade")
        .arg("rustup-init")
        .output()
    {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match Command::new("rustup").arg("update").output() {
        Ok(_) => (),
        Err(error) => match error.kind() {
            ErrorKind::NotFound => panic!("The rustup command is missing."),
            other_error => panic!("There was a problem: {:?}", other_error),
        },
    }
}

pub fn end() {
    // Environment (Unix)

    println!("Removing package manager artifacts.");

    match Command::new("brew").arg("cleanup").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Version control (Git)

    println!("Making sure all repos are clean.");

    let home_path = match dirs::home_dir() {
        Some(directory) => directory,
        None => panic!("Cannot find the home directory."),
    };

    let repos_path = home_path.join("repos");

    match fs::read_dir(&repos_path) {
        Ok(paths) => {
            let mut all_repos_clean = true;

            for path in paths {
                let path = match &path {
                    Ok(path) => path.path(),
                    Err(error) => panic!("There was a problem: {:?}", error),
                };

                let status_check = match Command::new("git")
                    .current_dir(&path)
                    .arg("status")
                    .arg("--porcelain")
                    .output()
                {
                    Ok(output) => output.stdout,
                    Err(error) => panic!("There was a problem: {:?}", error),
                };

                let unpushed_check = match Command::new("git")
                    .current_dir(&path)
                    .arg("log")
                    .arg("@{u}..")
                    .output()
                {
                    Ok(output) => output.stdout,
                    Err(error) => panic!("There was a problem: {:?}", error),
                };

                if !status_check.is_empty() || !unpushed_check.is_empty() {
                    all_repos_clean = false;
                    eprintln!("A dirty repo was found: {}", &path.display());
                }
            }

            if !all_repos_clean {
                panic!("Repos are dirty.");
            }
        }
        Err(error) => panic!("There was a problem: {:?}", error),
    };

    // Editor (Vim)

    println!("Quitting editor processes.");

    match Command::new("killall").arg("vim").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Application programming (JavaScript)

    println!("Quitting application programming processes.");

    match Command::new("killall").arg("node").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    // Systems programming (Rust)

    println!("Quitting systems programming processes.");

    match Command::new("killall").arg("rls").output() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::io::{BufReader, BufWriter};
use std::path::Path;

fn main() {
    generate(
        Path::new("data").join("adjectives.txt").as_path(),
        Path::new(&env::var("OUT_DIR").unwrap())
            .join("adjectives.rs")
            .as_path(),
    );
    generate(
        Path::new("data").join("nouns.txt").as_path(),
        Path::new(&env::var("OUT_DIR").unwrap())
            .join("nouns.rs")
            .as_path(),
    );
}

fn generate(src_path: &Path, dst_path: &Path) {
    let src = File::open(src_path).unwrap();
    let src = BufReader::new(src);
    let dst = File::create(dst_path).unwrap();
    let mut dst = BufWriter::new(dst);
    write!(dst, "[\n").unwrap();
    for word in src.lines() {
        write!(dst, "\"{}\",\n", &word.unwrap()).unwrap();
    }
    write!(dst, "];\n").unwrap();
}
<commit_msg> Remove stray semicolon<commit_after>use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::io::{BufReader, BufWriter};
use std::path::Path;

fn main() {
    generate(
        Path::new("data").join("adjectives.txt").as_path(),
        Path::new(&env::var("OUT_DIR").unwrap())
            .join("adjectives.rs")
            .as_path(),
    );
    generate(
        Path::new("data").join("nouns.txt").as_path(),
        Path::new(&env::var("OUT_DIR").unwrap())
            .join("nouns.rs")
            .as_path(),
    );
}

fn generate(src_path: &Path, dst_path: &Path) {
    let src = File::open(src_path).unwrap();
    let src = BufReader::new(src);
    let dst = File::create(dst_path).unwrap();
    let mut dst = BufWriter::new(dst);
    write!(dst, "[\n").unwrap();
    for word in src.lines() {
        write!(dst, "\"{}\",\n", &word.unwrap()).unwrap();
    }
    write!(dst, "]\n").unwrap();
}
<|endoftext|>


<commit_before>use file::tempfile;
use std::fs::File;
use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};
use std::mem::drop;

#[derive(Debug)]
enum SpooledInner {
    InMemory(Cursor<Vec<u8>>),
    OnDisk(File),
}

/// An object that behaves like a regular temporary file, but keeps data in
/// memory until it reaches a configured size, at which point the data is
/// written to a temporary file on disk, and further operations use the file
/// on disk.
#[derive(Debug)]
pub struct SpooledTempFile {
    max_size: usize,
    inner: SpooledInner,
}

/// Create a new spooled temporary file.
///
/// # Security
///
/// This variant is secure/reliable in the presence of a pathological temporary
/// file cleaner.
///
/// # Resource Leaking
///
/// The temporary file will be automatically removed by the OS when the last
/// handle to it is closed. This doesn't rely on Rust destructors being run, so
/// will (almost) never fail to clean up the temporary file.
///
/// # Examples
///
/// ```
/// # extern crate tempfile;
/// use tempfile::spooled_tempfile;
/// use std::io::{self, Write};
///
/// # fn main() {
/// #     if let Err(_) = run() {
/// #         ::std::process::exit(1);
/// #     }
/// # }
/// # fn run() -> Result<(), io::Error> {
/// let mut file = spooled_tempfile(15);
///
/// writeln!(file, "short line")?;
/// assert!(!file.is_rolled());
///
/// // as a result of this write call, the size of the data will exceed
/// // `max_size` (15), so it will be written to a temporary file on disk,
/// // and the in-memory buffer will be dropped
/// writeln!(file, "marvin gardens")?;
/// assert!(file.is_rolled());
///
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
    SpooledTempFile::new(max_size)
}

impl SpooledTempFile {
    pub fn new(max_size: usize) -> SpooledTempFile {
        SpooledTempFile {
            max_size: max_size,
            inner: SpooledInner::InMemory(Cursor::new(Vec::new())),
        }
    }

    /// Returns true if the file has been rolled over to disk.
    pub fn is_rolled(&self) -> bool {
        match self.inner {
            SpooledInner::InMemory(_) => false,
            SpooledInner::OnDisk(_) => true,
        }
    }

    /// Rolls over to a file on disk, regardless of current size. Does nothing
    /// if already rolled over.
    pub fn roll(&mut self) -> io::Result<()> {
        if !self.is_rolled() {
            let mut file = tempfile()?;
            if let SpooledInner::InMemory(ref mut cursor) = self.inner {
                file.write(cursor.get_ref())?;
                file.seek(SeekFrom::Start(cursor.position()))?;
                drop(cursor);
            }
            self.inner = SpooledInner::OnDisk(file);
        }
        Ok(())
    }

    pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
        if size as usize > self.max_size {
            self.roll()?; // does nothing if already rolled over
        }
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => {
                cursor.get_mut().resize(size as usize, 0);
                Ok(())
            }
            SpooledInner::OnDisk(ref mut file) => file.set_len(size),
        }
    }
}

impl Read for SpooledTempFile {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.read(buf),
            SpooledInner::OnDisk(ref mut file) => file.read(buf),
        }
    }
}

impl Write for SpooledTempFile {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        // roll over to file if necessary
        let mut rolling = false;
        if let SpooledInner::InMemory(ref mut cursor) = self.inner {
            rolling = cursor.position() as usize + buf.len() > self.max_size;
        }
        if rolling {
            self.roll()?;
        }

        // write the bytes
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.write(buf),
            SpooledInner::OnDisk(ref mut file) => file.write(buf),
        }
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.flush(),
            SpooledInner::OnDisk(ref mut file) => file.flush(),
        }
    }
}

impl Seek for SpooledTempFile {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.seek(pos),
            SpooledInner::OnDisk(ref mut file) => file.seek(pos),
        }
    }
}
<commit_msg>fix: use write_all instead of write when writing spooled data to file<commit_after>use file::tempfile;
use std::fs::File;
use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};

#[derive(Debug)]
enum SpooledInner {
    InMemory(Cursor<Vec<u8>>),
    OnDisk(File),
}

/// An object that behaves like a regular temporary file, but keeps data in
/// memory until it reaches a configured size, at which point the data is
/// written to a temporary file on disk, and further operations use the file
/// on disk.
#[derive(Debug)]
pub struct SpooledTempFile {
    max_size: usize,
    inner: SpooledInner,
}

/// Create a new spooled temporary file.
///
/// # Security
///
/// This variant is secure/reliable in the presence of a pathological temporary
/// file cleaner.
///
/// # Resource Leaking
///
/// The temporary file will be automatically removed by the OS when the last
/// handle to it is closed. This doesn't rely on Rust destructors being run, so
/// will (almost) never fail to clean up the temporary file.
///
/// # Examples
///
/// ```
/// # extern crate tempfile;
/// use tempfile::spooled_tempfile;
/// use std::io::{self, Write};
///
/// # fn main() {
/// #     if let Err(_) = run() {
/// #         ::std::process::exit(1);
/// #     }
/// # }
/// # fn run() -> Result<(), io::Error> {
/// let mut file = spooled_tempfile(15);
///
/// writeln!(file, "short line")?;
/// assert!(!file.is_rolled());
///
/// // as a result of this write call, the size of the data will exceed
/// // `max_size` (15), so it will be written to a temporary file on disk,
/// // and the in-memory buffer will be dropped
/// writeln!(file, "marvin gardens")?;
/// assert!(file.is_rolled());
///
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
    SpooledTempFile::new(max_size)
}

impl SpooledTempFile {
    pub fn new(max_size: usize) -> SpooledTempFile {
        SpooledTempFile {
            max_size: max_size,
            inner: SpooledInner::InMemory(Cursor::new(Vec::new())),
        }
    }

    /// Returns true if the file has been rolled over to disk.
    pub fn is_rolled(&self) -> bool {
        match self.inner {
            SpooledInner::InMemory(_) => false,
            SpooledInner::OnDisk(_) => true,
        }
    }

    /// Rolls over to a file on disk, regardless of current size. Does nothing
    /// if already rolled over.
    pub fn roll(&mut self) -> io::Result<()> {
        if !self.is_rolled() {
            let mut file = tempfile()?;
            if let SpooledInner::InMemory(ref mut cursor) = self.inner {
                file.write_all(cursor.get_ref())?;
                file.seek(SeekFrom::Start(cursor.position()))?;
            }
            self.inner = SpooledInner::OnDisk(file);
        }
        Ok(())
    }

    pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
        if size as usize > self.max_size {
            self.roll()?; // does nothing if already rolled over
        }
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => {
                cursor.get_mut().resize(size as usize, 0);
                Ok(())
            }
            SpooledInner::OnDisk(ref mut file) => file.set_len(size),
        }
    }
}

impl Read for SpooledTempFile {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.read(buf),
            SpooledInner::OnDisk(ref mut file) => file.read(buf),
        }
    }
}

impl Write for SpooledTempFile {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        // roll over to file if necessary
        let mut rolling = false;
        if let SpooledInner::InMemory(ref mut cursor) = self.inner {
            rolling = cursor.position() as usize + buf.len() > self.max_size;
        }
        if rolling {
            self.roll()?;
        }

        // write the bytes
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.write(buf),
            SpooledInner::OnDisk(ref mut file) => file.write(buf),
        }
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.flush(),
            SpooledInner::OnDisk(ref mut file) => file.flush(),
        }
    }
}

impl Seek for SpooledTempFile {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        match self.inner {
            SpooledInner::InMemory(ref mut cursor) => cursor.seek(pos),
            SpooledInner::OnDisk(ref mut file) => file.seek(pos),
        }
    }
}
<|endoftext|>


<commit_before>use std::{mem, str};
use std::path::{Path, PathBuf};
use std::io::Write;
use std::sync::Arc;
use rocksdb::{self, Writable};
use byteorder::{BigEndian, WriteBytesExt};
use utils::*;
use nodrop::NoDrop;

struct U16BeSuffixTransform;

impl rocksdb::SliceTransform for U16BeSuffixTransform {
    fn transform<'a>(&mut self, key: &'a [u8]) -> &'a [u8] {
        &key[..2]
    }

    fn in_domain(&mut self, key: &[u8]) -> bool {
        key.len() >= 2
    }
}

pub struct StorageManager {
    path: PathBuf,
    db: Arc<rocksdb::DB>,
}

// TODO: merge operator would be a huge win
// This implementation goes to a lot of trouble to allow
// safe-ish iterators that don't require a lifetime attached to them.
pub struct Storage {
    db: Arc<rocksdb::DB>,
    iterators_handle: Arc<()>,
    num: u16,
}

pub struct StorageIterator {
    db: Arc<rocksdb::DB>,
    snapshot: NoDrop<Box<rocksdb::rocksdb::Snapshot<'static>>>,
    iterator: NoDrop<Box<rocksdb::rocksdb::DBIterator<'static>>>,
    iterators_handle: Arc<()>,
    key_prefix: [u8; 2],
    first: bool,
}

unsafe impl Send for StorageIterator {}

impl StorageManager {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<StorageManager, GenericError> {
        let mut opts = rocksdb::Options::new();
        opts.create_if_missing(true);
        opts.set_prefix_extractor("U16BeSuffixTransform", Box::new(U16BeSuffixTransform))
            .unwrap();
        opts.compression_per_level(
            &[
                rocksdb::DBCompressionType::DBNo,
                rocksdb::DBCompressionType::DBNo,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
            ],
        );
        opts.set_write_buffer_size(32 * 1024 * 1024);
        opts.set_max_bytes_for_level_base(4 * 32 * 1024 * 1024);
        opts.set_max_write_buffer_number(4);
        opts.set_max_background_flushes(2);
        opts.set_max_background_compactions(2);
        let mut block_opts = rocksdb::BlockBasedOptions::new();
        block_opts.set_bloom_filter(10, false);
        block_opts.set_lru_cache(4 * 32 * 1024 * 1024);
        opts.set_block_based_table_factory(&block_opts);
        // TODO: Rocksdb is complicated, we might want to tune some more options
        let db = rocksdb::DB::open(opts, path.as_ref().to_str().unwrap()).unwrap();
        Ok(
            StorageManager {
                path: path.as_ref().into(),
                db: Arc::new(db),
            },
        )
    }

    pub fn open(&self, db_num: u16, _create: bool) -> Result<Storage, GenericError> {
        Ok(
            Storage {
                db: self.db.clone(),
                num: db_num,
                iterators_handle: Arc::new(()),
            },
        )
    }
}

impl Drop for StorageManager {
    fn drop(&mut self) {
        let sc = Arc::strong_count(&self.db);
        let wc = Arc::weak_count(&self.db);
        assert!(wc == 0);
        assert!(sc == 1, "{} pending databases", sc - 1);
    }
}

impl Storage {
    #[inline]
    fn build_key<'a>(&self, buffer: &'a mut [u8], key: &[u8]) -> &'a [u8] {
        (&mut buffer[..2]).write_u16::<BigEndian>(self.num).unwrap();
        (&mut buffer[2..]).write_all(key).unwrap();
        &buffer[.. 2 + key.len()]
    }

    pub fn get<R, F: FnOnce(&[u8]) -> R>(&self, key: &[u8], callback: F) -> Option<R> {
        let mut buffer = [0u8; 512];
        let buffer = self.build_key(&mut buffer, key);
        let r = self.db.get(buffer).unwrap();
        trace!("get {:?} ({:?} bytes)", str::from_utf8(key), r.as_ref().map(|x| x.len()));
        r.map(|r| callback(&*r))
    }

    pub fn iterator(&self) -> StorageIterator {
        let mut key_prefix = [0u8; 2];
        (&mut key_prefix[..]).write_u16::<BigEndian>(self.num).unwrap();
        unsafe {
            let snapshot = NoDrop::new(Box::new(self.db.snapshot()));
            let mut iterator = NoDrop::new(Box::new(snapshot.iter()));
            iterator.seek(rocksdb::SeekKey::Key(&key_prefix[..]));
            let result = StorageIterator {
                db: self.db.clone(),
                snapshot: mem::transmute_copy(&snapshot),
                iterator: mem::transmute_copy(&iterator),
                iterators_handle: self.iterators_handle.clone(),
                key_prefix: key_prefix,
                first: true,
            };
            result
        }
    }

    pub fn get_vec(&self, key: &[u8]) -> Option<Vec<u8>> {
        self.get(key, |v| v.to_owned())
    }

    pub fn set(&self, key: &[u8], value: &[u8]) {
        trace!("set {:?} ({} bytes)", str::from_utf8(key), value.len());
        let mut buffer = [0u8; 512];
        let buffer = self.build_key(&mut buffer, key);
        self.db.put(buffer, value).unwrap();
    }

    pub fn del(&self, key: &[u8]) {
        trace!("del {:?}", str::from_utf8(key));
        let mut buffer = [0u8; 512];
        let buffer = self.build_key(&mut buffer, key);
        self.db.delete(buffer).unwrap()
    }

    pub fn clear(&self) {
        trace!("clear");
        self.check_pending_iters();
        let mut from = [0u8; 2];
        let mut to = [0u8; 2];
        (&mut from[..]).write_u16::<BigEndian>(self.num).unwrap();
        (&mut to[..]).write_u16::<BigEndian>(self.num + 1).unwrap();
        self.db.delete_file_in_range(&from[..], &to[..]).unwrap();
        for (key, _) in self.iterator().iter() {
            self.db.delete(key).unwrap();
        }
    }

    pub fn sync(&self) {
        debug!("sync");
        self.db.sync_wal().unwrap();
    }

    fn check_pending_iters(&self) {
        let sc = Arc::strong_count(&self.iterators_handle);
        let wc = Arc::weak_count(&self.iterators_handle);
        assert!(wc == 0);
        assert!(sc == 1, "{} pending databases", sc - 1);
    }
}

impl Drop for Storage {
    fn drop(&mut self) {
        self.check_pending_iters();
    }
}

impl StorageIterator {
    pub fn iter<'a>(&'a mut self) -> StorageIter<'a> {
        StorageIter { it: self }
    }
}

pub struct StorageIter<'a> {
    it: &'a mut StorageIterator,
}

impl<'a> Iterator for StorageIter<'a> {
    type Item = (&'a [u8], &'a [u8]);
    fn next(&mut self) -> Option<Self::Item> {
        if self.it.first {
            self.it.first = false;
        } else {
            self.it.iterator.next();
        }
        if self.it.iterator.valid() && self.it.iterator.key()[..2] == self.it.key_prefix[..] {
            unsafe {
                // safe as slices are valid until the next call to next
                Some(
                    (mem::transmute(&self.it.iterator.key()[2..]),
                     mem::transmute(self.it.iterator.value())),
                )
            }
        } else {
            None
        }
    }
}

impl Drop for StorageIterator {
    fn drop(&mut self) {
        unsafe {
            let iterator: NoDrop<Box<rocksdb::rocksdb::DBIterator>> = mem::transmute_copy(&self.iterator,);
            let snapshot: NoDrop<Box<rocksdb::rocksdb::Snapshot>> = mem::transmute_copy(&self.snapshot,);
            drop(iterator.into_inner());
            drop(snapshot.into_inner());
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::Path;
    use std::fs;

    #[test]
    fn test_simple() {
        let _ = fs::remove_dir_all("t/test_simple");
        let sm = StorageManager::new("t/test_simple").unwrap();
        let storage = sm.open(1, true).unwrap();
        assert_eq!(storage.get_vec(b"sample"), None);
        storage.set(b"sample", b"sample_value");
        assert_eq!(storage.get_vec(b"sample").unwrap(), b"sample_value");
        storage.del(b"sample");
        assert_eq!(storage.get_vec(b"sample"), None);
        drop(storage);
    }

    #[test]
    fn test_iter2() {
        let _ = fs::remove_dir_all("t/test_iter2");
        let sm = StorageManager::new("t/test_iter2").unwrap();
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            storage.set(b"1", i.to_string().as_bytes());
            storage.set(b"2", i.to_string().as_bytes());
            storage.set(b"3", i.to_string().as_bytes());
        }
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            let results: Vec<Vec<u8>> = storage.iterator().iter().map(|(k, v)| v.into()).collect();
            assert_eq!(
                results,
                &[i.to_string().as_bytes(), i.to_string().as_bytes(), i.to_string().as_bytes()]
            );
        }
    }

    #[test]
    fn test_clear() {
        let _ = fs::remove_dir_all("t/test_clear");
        let sm = StorageManager::new("t/test_clear").unwrap();
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            storage.set(b"1", i.to_string().as_bytes());
            storage.set(b"2", i.to_string().as_bytes());
            storage.set(b"3", i.to_string().as_bytes());
        }
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            let results: Vec<Vec<u8>> = storage.iterator().iter().map(|(k, v)| v.into()).collect();
            assert_eq!(
                results,
                &[i.to_string().as_bytes(), i.to_string().as_bytes(), i.to_string().as_bytes()]
            );
            storage.clear();
            let results: Vec<Vec<u8>> = storage.iterator().iter().map(|(k, v)| v.into()).collect();
            assert_eq!(results.len(), 0);
        }
    }

    #[test]
    fn test_open_all() {
        let _ = fs::remove_dir_all("t/test_open_all");
        let sm = StorageManager::new("t/test_open_all").unwrap();
        sm.open(1, true).unwrap();
        sm.open(2, true).unwrap();
        sm.open(3, true).unwrap();
        sm.open(1, false).unwrap();
        sm.open(2, false).unwrap();
        sm.open(3, false).unwrap();
    }

}
<commit_msg>fix storage.clear<commit_after>use std::{mem, str};
use std::path::{Path, PathBuf};
use std::io::Write;
use std::sync::Arc;
use rocksdb::{self, Writable};
use byteorder::{BigEndian, WriteBytesExt};
use utils::*;
use nodrop::NoDrop;

struct U16BeSuffixTransform;

impl rocksdb::SliceTransform for U16BeSuffixTransform {
    fn transform<'a>(&mut self, key: &'a [u8]) -> &'a [u8] {
        &key[..2]
    }

    fn in_domain(&mut self, key: &[u8]) -> bool {
        key.len() >= 2
    }
}

pub struct StorageManager {
    path: PathBuf,
    db: Arc<rocksdb::DB>,
}

// TODO: merge operator would be a huge win
// This implementation goes to a lot of trouble to allow
// safe-ish iterators that don't require a lifetime attached to them.
pub struct Storage {
    db: Arc<rocksdb::DB>,
    iterators_handle: Arc<()>,
    num: u16,
}

pub struct StorageIterator {
    db: Arc<rocksdb::DB>,
    snapshot: NoDrop<Box<rocksdb::rocksdb::Snapshot<'static>>>,
    iterator: NoDrop<Box<rocksdb::rocksdb::DBIterator<'static>>>,
    iterators_handle: Arc<()>,
    key_prefix: [u8; 2],
    first: bool,
}

unsafe impl Send for StorageIterator {}

impl StorageManager {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<StorageManager, GenericError> {
        let mut opts = rocksdb::Options::new();
        opts.create_if_missing(true);
        opts.set_prefix_extractor("U16BeSuffixTransform", Box::new(U16BeSuffixTransform))
            .unwrap();
        opts.compression_per_level(
            &[
                rocksdb::DBCompressionType::DBNo,
                rocksdb::DBCompressionType::DBNo,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
                rocksdb::DBCompressionType::DBLz4,
            ],
        );
        opts.set_write_buffer_size(32 * 1024 * 1024);
        opts.set_max_bytes_for_level_base(4 * 32 * 1024 * 1024);
        opts.set_max_write_buffer_number(4);
        opts.set_max_background_flushes(2);
        opts.set_max_background_compactions(2);
        let mut block_opts = rocksdb::BlockBasedOptions::new();
        block_opts.set_bloom_filter(10, false);
        block_opts.set_lru_cache(4 * 32 * 1024 * 1024);
        opts.set_block_based_table_factory(&block_opts);
        // TODO: Rocksdb is complicated, we might want to tune some more options
        let db = rocksdb::DB::open(opts, path.as_ref().to_str().unwrap()).unwrap();
        Ok(
            StorageManager {
                path: path.as_ref().into(),
                db: Arc::new(db),
            },
        )
    }

    pub fn open(&self, db_num: u16, _create: bool) -> Result<Storage, GenericError> {
        Ok(
            Storage {
                db: self.db.clone(),
                num: db_num,
                iterators_handle: Arc::new(()),
            },
        )
    }
}

impl Drop for StorageManager {
    fn drop(&mut self) {
        let sc = Arc::strong_count(&self.db);
        let wc = Arc::weak_count(&self.db);
        assert!(wc == 0);
        assert!(sc == 1, "{} pending databases", sc - 1);
    }
}

impl Storage {
    #[inline]
    fn build_key<'a>(&self, buffer: &'a mut [u8], key: &[u8]) -> &'a [u8] {
        (&mut buffer[..2]).write_u16::<BigEndian>(self.num).unwrap();
        (&mut buffer[2..]).write_all(key).unwrap();
        &buffer[.. 2 + key.len()]
    }

    pub fn get<R, F: FnOnce(&[u8]) -> R>(&self, key: &[u8], callback: F) -> Option<R> {
        let mut buffer = [0u8; 512];
        let buffer = self.build_key(&mut buffer, key);
        let r = self.db.get(buffer).unwrap();
        trace!("get {:?} ({:?} bytes)", str::from_utf8(key), r.as_ref().map(|x| x.len()));
        r.map(|r| callback(&*r))
    }

    pub fn iterator(&self) -> StorageIterator {
        let mut key_prefix = [0u8; 2];
        (&mut key_prefix[..]).write_u16::<BigEndian>(self.num).unwrap();
        unsafe {
            let snapshot = NoDrop::new(Box::new(self.db.snapshot()));
            let mut iterator = NoDrop::new(Box::new(snapshot.iter()));
            iterator.seek(rocksdb::SeekKey::Key(&key_prefix[..]));
            let result = StorageIterator {
                db: self.db.clone(),
                snapshot: mem::transmute_copy(&snapshot),
                iterator: mem::transmute_copy(&iterator),
                iterators_handle: self.iterators_handle.clone(),
                key_prefix: key_prefix,
                first: true,
            };
            result
        }
    }

    pub fn get_vec(&self, key: &[u8]) -> Option<Vec<u8>> {
        self.get(key, |v| v.to_owned())
    }

    pub fn set(&self, key: &[u8], value: &[u8]) {
        trace!("set {:?} ({} bytes)", str::from_utf8(key), value.len());
        let mut buffer = [0u8; 512];
        let buffer = self.build_key(&mut buffer, key);
        self.db.put(buffer, value).unwrap();
    }

    pub fn del(&self, key: &[u8]) {
        trace!("del {:?}", str::from_utf8(key));
        let mut buffer = [0u8; 512];
        let buffer = self.build_key(&mut buffer, key);
        self.db.delete(buffer).unwrap()
    }

    pub fn clear(&self) {
        trace!("clear");
        self.check_pending_iters();
        let mut from = [0u8; 2];
        let mut to = [0u8; 2];
        (&mut from[..]).write_u16::<BigEndian>(self.num).unwrap();
        (&mut to[..]).write_u16::<BigEndian>(self.num + 1).unwrap();
        self.db.delete_file_in_range(&from[..], &to[..]).unwrap();
        for (key, _) in self.iterator().iter() {
            self.del(key);
        }
    }

    pub fn sync(&self) {
        debug!("sync");
        self.db.sync_wal().unwrap();
    }

    fn check_pending_iters(&self) {
        let sc = Arc::strong_count(&self.iterators_handle);
        let wc = Arc::weak_count(&self.iterators_handle);
        assert!(wc == 0);
        assert!(sc == 1, "{} pending databases", sc - 1);
    }
}

impl Drop for Storage {
    fn drop(&mut self) {
        self.check_pending_iters();
    }
}

impl StorageIterator {
    pub fn iter<'a>(&'a mut self) -> StorageIter<'a> {
        StorageIter { it: self }
    }
}

pub struct StorageIter<'a> {
    it: &'a mut StorageIterator,
}

impl<'a> Iterator for StorageIter<'a> {
    type Item = (&'a [u8], &'a [u8]);
    fn next(&mut self) -> Option<Self::Item> {
        if self.it.first {
            self.it.first = false;
        } else {
            self.it.iterator.next();
        }
        if self.it.iterator.valid() && self.it.iterator.key()[..2] == self.it.key_prefix[..] {
            unsafe {
                // safe as slices are valid until the next call to next
                Some(
                    (mem::transmute(&self.it.iterator.key()[2..]),
                     mem::transmute(self.it.iterator.value())),
                )
            }
        } else {
            None
        }
    }
}

impl Drop for StorageIterator {
    fn drop(&mut self) {
        unsafe {
            let iterator: NoDrop<Box<rocksdb::rocksdb::DBIterator>> = mem::transmute_copy(&self.iterator,);
            let snapshot: NoDrop<Box<rocksdb::rocksdb::Snapshot>> = mem::transmute_copy(&self.snapshot,);
            drop(iterator.into_inner());
            drop(snapshot.into_inner());
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::Path;
    use std::fs;

    #[test]
    fn test_simple() {
        let _ = fs::remove_dir_all("t/test_simple");
        let sm = StorageManager::new("t/test_simple").unwrap();
        let storage = sm.open(1, true).unwrap();
        assert_eq!(storage.get_vec(b"sample"), None);
        storage.set(b"sample", b"sample_value");
        assert_eq!(storage.get_vec(b"sample").unwrap(), b"sample_value");
        storage.del(b"sample");
        assert_eq!(storage.get_vec(b"sample"), None);
        drop(storage);
    }

    #[test]
    fn test_iter2() {
        let _ = fs::remove_dir_all("t/test_iter2");
        let sm = StorageManager::new("t/test_iter2").unwrap();
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            storage.set(b"1", i.to_string().as_bytes());
            storage.set(b"2", i.to_string().as_bytes());
            storage.set(b"3", i.to_string().as_bytes());
        }
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            let results: Vec<Vec<u8>> = storage.iterator().iter().map(|(k, v)| v.into()).collect();
            assert_eq!(
                results,
                &[i.to_string().as_bytes(), i.to_string().as_bytes(), i.to_string().as_bytes()]
            );
        }
    }

    #[test]
    fn test_clear() {
        let _ = fs::remove_dir_all("t/test_clear");
        let sm = StorageManager::new("t/test_clear").unwrap();
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            storage.set(b"1", i.to_string().as_bytes());
            storage.set(b"2", i.to_string().as_bytes());
            storage.set(b"3", i.to_string().as_bytes());
        }
        for &i in &[0, 1, 2] {
            let storage = sm.open(i, true).unwrap();
            let results: Vec<Vec<u8>> = storage.iterator().iter().map(|(k, v)| v.into()).collect();
            assert_eq!(
                results,
                &[i.to_string().as_bytes(), i.to_string().as_bytes(), i.to_string().as_bytes()]
            );
            storage.clear();
            let results: Vec<Vec<u8>> = storage.iterator().iter().map(|(k, v)| v.into()).collect();
            assert_eq!(results.len(), 0);
        }
    }

    #[test]
    fn test_open_all() {
        let _ = fs::remove_dir_all("t/test_open_all");
        let sm = StorageManager::new("t/test_open_all").unwrap();
        sm.open(1, true).unwrap();
        sm.open(2, true).unwrap();
        sm.open(3, true).unwrap();
        sm.open(1, false).unwrap();
        sm.open(2, false).unwrap();
        sm.open(3, false).unwrap();
    }

}
<|endoftext|>


<commit_before>pub use luminance::pixel::{Depth32F, R32F, RGBA32F};
pub use luminance::texture::{Dim2, Flat, MagFilter, MinFilter, Sampler, Texture, Unit, Wrap};
use image;
use std::ops::Deref;
use std::path::Path;

use resource::{Load, LoadError, Reload, ResCache, Result};

// Common texture aliases.
pub type TextureRGBA32F = Texture<Flat, Dim2, RGBA32F>;
pub type TextureDepth32F = Texture<Flat, Dim2, Depth32F>;

/// Load an RGBA texture from an image at a path.
///
/// The `linearizer` argument is an option that gives the factor to apply to linearize if needed. Pass
/// `None` if the texture is already linearized.
pub fn load_rgba_texture<P, L>(path: P, sampler: &Sampler, linearizer: L) -> Result<TextureRGBA32F> where P: AsRef<Path>, L: Into<Option<f32>> {
  info!("loading texture image: \x1b[35m{:?}", path.as_ref());

  let img = image::open(path).map_err(|e| LoadError::ConversionFailed(format!("{:?}", e)))?.flipv().to_rgba();
  let (w, h) = img.dimensions();
  let linearizer = linearizer.into();
  let raw: Vec<f32> = img.into_raw().into_iter().map(|x| {
    let y = x as f32 / 255.;
    linearizer.map_or(y, |factor| y.powf(1. / factor))
  }).collect();

  let tex = Texture::new([w, h], 0, sampler).map_err(|e| LoadError::ConversionFailed(format!("{:?}", e)))?;
  tex.upload_raw(false, &raw);

  Ok(tex)
}

/// Save an RGBA image on disk.
pub fn save_rgba_texture<P>(texture: &TextureRGBA32F, path: P) where P: AsRef<Path> {
  info!("saving texture image to: \x1b[35m{:?}", path.as_ref());

  let texels = texture.get_raw_texels();
  let [w, h] = texture.size();
  let mut output = Vec::with_capacity((w * h) as usize);

  for texel in &texels {
    output.push((texel * 255.) as u8);
  }

  let _ = image::save_buffer(path, &output, w, h, image::ColorType::RGBA(8));
}

pub struct TextureImage {
  pub texture: TextureRGBA32F,
  sampler: Sampler,
  linearizer: Option<f32>,
}

impl Deref for TextureImage {
  type Target = TextureRGBA32F;

  fn deref(&self) -> &Self::Target {
    &self.texture
  }
}

impl Load for TextureImage {
  type Args = (Sampler, Option<f32>);

  const TY_STR: &'static str = "textures";

  fn load<P>(path: P, _: &mut ResCache, (sampler, linearizer): Self::Args) -> Result<Self> where P: AsRef<Path> {
    load_rgba_texture(path, &sampler, linearizer)
      .map(|tex| TextureImage {
        texture: tex,
        sampler: sampler,
        linearizer: linearizer
      })
  }
}

impl Reload for TextureImage {
  fn reload_args(&self) -> Self::Args {
    (self.sampler, self.linearizer)
  }
}
<commit_msg>Made all TextureImage’s fields pub.<commit_after>pub use luminance::pixel::{Depth32F, R32F, RGBA32F};
pub use luminance::texture::{Dim2, Flat, MagFilter, MinFilter, Sampler, Texture, Unit, Wrap};
use image;
use std::ops::Deref;
use std::path::Path;

use resource::{Load, LoadError, Reload, ResCache, Result};

// Common texture aliases.
pub type TextureRGBA32F = Texture<Flat, Dim2, RGBA32F>;
pub type TextureDepth32F = Texture<Flat, Dim2, Depth32F>;

/// Load an RGBA texture from an image at a path.
///
/// The `linearizer` argument is an option that gives the factor to apply to linearize if needed. Pass
/// `None` if the texture is already linearized.
pub fn load_rgba_texture<P, L>(path: P, sampler: &Sampler, linearizer: L) -> Result<TextureRGBA32F> where P: AsRef<Path>, L: Into<Option<f32>> {
  info!("loading texture image: \x1b[35m{:?}", path.as_ref());

  let img = image::open(path).map_err(|e| LoadError::ConversionFailed(format!("{:?}", e)))?.flipv().to_rgba();
  let (w, h) = img.dimensions();
  let linearizer = linearizer.into();
  let raw: Vec<f32> = img.into_raw().into_iter().map(|x| {
    let y = x as f32 / 255.;
    linearizer.map_or(y, |factor| y.powf(1. / factor))
  }).collect();

  let tex = Texture::new([w, h], 0, sampler).map_err(|e| LoadError::ConversionFailed(format!("{:?}", e)))?;
  tex.upload_raw(false, &raw);

  Ok(tex)
}

/// Save an RGBA image on disk.
pub fn save_rgba_texture<P>(texture: &TextureRGBA32F, path: P) where P: AsRef<Path> {
  info!("saving texture image to: \x1b[35m{:?}", path.as_ref());

  let texels = texture.get_raw_texels();
  let [w, h] = texture.size();
  let mut output = Vec::with_capacity((w * h) as usize);

  for texel in &texels {
    output.push((texel * 255.) as u8);
  }

  let _ = image::save_buffer(path, &output, w, h, image::ColorType::RGBA(8));
}

pub struct TextureImage {
  pub texture: TextureRGBA32F,
  pub sampler: Sampler,
  pub linearizer: Option<f32>,
}

impl Deref for TextureImage {
  type Target = TextureRGBA32F;

  fn deref(&self) -> &Self::Target {
    &self.texture
  }
}

impl Load for TextureImage {
  type Args = (Sampler, Option<f32>);

  const TY_STR: &'static str = "textures";

  fn load<P>(path: P, _: &mut ResCache, (sampler, linearizer): Self::Args) -> Result<Self> where P: AsRef<Path> {
    load_rgba_texture(path, &sampler, linearizer)
      .map(|tex| TextureImage {
        texture: tex,
        sampler: sampler,
        linearizer: linearizer
      })
  }
}

impl Reload for TextureImage {
  fn reload_args(&self) -> Self::Args {
    (self.sampler, self.linearizer)
  }
}
<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::io;
use std::net::{UdpSocket, SocketAddr, ToSocketAddrs};

pub struct Transit<T> {
    socket: UdpSocket,
    packet_type: PhantomData<T>,
}

pub trait FromTransit<T> {
    fn from(T) -> io::Result<Self>;
}

impl<T: for<'a> FromTransit<&'a [u8]> + Into<Vec<u8>>> Transit<T> {
    pub fn new<A>(addr: A) -> io::Result<Transit<T>> where A: ToSocketAddrs {
        let socket = try!(UdpSocket::bind(addr));
        Ok(Transit {
            socket: socket,
            packet_type: PhantomData,
        })
    }

    pub fn recv_from(&self) -> io::Result<(T, SocketAddr)> {
        let mut buf = [0; 1024];
        let (n, addr) = try!(self.socket.recv_from(&mut buf));
        assert!(n <= 1024);
        let data = try!(T::from(&buf));
        Ok((data, addr))
    }

    pub fn send_to<A>(&self, pkt: T, addr: A) -> io::Result<()> where A: ToSocketAddrs {
        let buf = pkt.into();
        try!(self.socket.send_to(buf.as_slice(), addr));
        Ok(())
    }

    pub fn local_addr(&self)  -> io::Result<SocketAddr> {
        self.socket.local_addr()
    }
}

#[cfg(test)]
mod test {
    use udp::*;
    use std::io::{self, Error, ErrorKind};

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Test {
        ten: u8,
    }

    impl Into<Vec<u8>> for Test {
        fn into(self) -> Vec<u8> {
            let mut vec  = Vec::new();
            vec.push(self.ten);
            vec
        }
    }

    impl<'a> FromTransit<&'a [u8]> for Test {
        fn from(buf: &'a [u8]) -> io::Result<Test> {
            if buf[0] != 10 {
                Err(Error::new(ErrorKind::InvalidData, "failed to serialize"))
            } else {
                Ok(Test { ten: buf[0] })
            }
        }
    }

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Another {
        data: u8,
    }

    impl Into<Vec<u8>> for Another {
        fn into(self) -> Vec<u8> {
            let mut vec  = Vec::new();
            vec.push(self.data);
            vec
        }
    }

    impl<'a> FromTransit<&'a [u8]> for Another {
        fn from(buf: &'a [u8]) -> io::Result<Another> {
            Ok(Another { data: buf[0] })
        }
    }

    #[test]
    fn test_send_recv() {
        let addr1 = "127.0.0.1:61001";
        let addr2 = "127.0.0.1:61002";
        let transit1: Transit<Test> = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let test = Test { ten: 10 };

        let res = transit2.send_to(test, addr1);
        assert!(res.is_ok());
        let res = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, test);
    }

    #[test]
    fn test_packet_type() {
        let addr1 = "127.0.0.1:62001";
        let addr2 = "127.0.0.1:62002";
        let transit1: Transit<Another> = Transit::new(addr1).unwrap();
        let transit2: Transit<Test> = Transit::new(addr2).unwrap();
        let test = Another { data: 27 };

        let res = transit1.send_to(test, addr2);
        assert!(res.is_ok());
        let res = transit2.recv_from();
        assert!(res.is_err());
    }
}
<commit_msg>Add some documentation<commit_after>use std::marker::PhantomData;
use std::io;
use std::net::{UdpSocket, SocketAddr, ToSocketAddrs};

pub struct Transit<T> {
    socket: UdpSocket,
    packet_type: PhantomData<T>,
}

/// The From trait in the standard library does not allow for errors to occur. This is useful
/// particularly when converting a type from a byte array received from the network that may or may
/// not be of the particular type.
pub trait FromTransit<T> {
    fn from(T) -> io::Result<Self>;
}

impl<T: for<'a> FromTransit<&'a [u8]> + Into<Vec<u8>>> Transit<T> {
    pub fn new<A>(addr: A) -> io::Result<Transit<T>> where A: ToSocketAddrs {
        let socket = try!(UdpSocket::bind(addr));
        Ok(Transit {
            socket: socket,
            packet_type: PhantomData,
        })
    }

    /// On success, this function returns the type deserialized using the FromTransit trait
    /// implementation. The trait implementation should be able to detect whether or not the buffer
    /// contains a valid UDP message and emit an error appropriately.
    pub fn recv_from(&self) -> io::Result<(T, SocketAddr)> {
        let mut buf = [0; 1024];
        let (n, addr) = try!(self.socket.recv_from(&mut buf));
        assert!(n <= 1024);
        let data = try!(T::from(&buf));
        Ok((data, addr))
    }

    /// Transforms the packet into a byte array and sends it to the associated address.
    pub fn send_to<A>(&self, pkt: T, addr: A) -> io::Result<()> where A: ToSocketAddrs {
        let buf = pkt.into();
        try!(self.socket.send_to(buf.as_slice(), addr));
        Ok(())
    }

    pub fn local_addr(&self)  -> io::Result<SocketAddr> {
        self.socket.local_addr()
    }
}

#[cfg(test)]
mod test {
    use udp::*;
    use std::io::{self, Error, ErrorKind};

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Test {
        ten: u8,
    }

    impl Into<Vec<u8>> for Test {
        fn into(self) -> Vec<u8> {
            let mut vec  = Vec::new();
            vec.push(self.ten);
            vec
        }
    }

    impl<'a> FromTransit<&'a [u8]> for Test {
        fn from(buf: &'a [u8]) -> io::Result<Test> {
            if buf[0] != 10 {
                Err(Error::new(ErrorKind::InvalidData, "failed to serialize"))
            } else {
                Ok(Test { ten: buf[0] })
            }
        }
    }

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Another {
        data: u8,
    }

    impl Into<Vec<u8>> for Another {
        fn into(self) -> Vec<u8> {
            let mut vec  = Vec::new();
            vec.push(self.data);
            vec
        }
    }

    impl<'a> FromTransit<&'a [u8]> for Another {
        fn from(buf: &'a [u8]) -> io::Result<Another> {
            Ok(Another { data: buf[0] })
        }
    }

    #[test]
    fn test_send_recv() {
        let addr1 = "127.0.0.1:61001";
        let addr2 = "127.0.0.1:61002";
        let transit1: Transit<Test> = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let test = Test { ten: 10 };

        let res = transit2.send_to(test, addr1);
        assert!(res.is_ok());
        let res = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, test);
    }

    #[test]
    fn test_packet_type() {
        let addr1 = "127.0.0.1:62001";
        let addr2 = "127.0.0.1:62002";
        let transit1: Transit<Another> = Transit::new(addr1).unwrap();
        let transit2: Transit<Test> = Transit::new(addr2).unwrap();
        let test = Another { data: 27 };

        let res = transit1.send_to(test, addr2);
        assert!(res.is_ok());
        let res = transit2.recv_from();
        assert!(res.is_err());
    }
}
<|endoftext|>