<commit_before>extern crate regex;

use super::{Registry, Module, HandleData};
use lib::{self, Error, Request, db};
use lib::downloader::Filename;
use std::fs::remove_file;
use std::fs::remove_dir_all;
use std::fs::create_dir;
use std::path::Path;

use CODE_IN_PROGRESS;
use CONFIG;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

macro_rules! condition(
    ($e:expr,$y:expr) => (match $y { true => {$e;}, false => ()} );
);

const YT_VIDEO_URL: &'static str = "https://www.youtube.com/watch?v=";

lazy_static! {
// https://regex101.com/r/lZ6lC1/3
// we need to remove the / escaping!
    pub static ref REGEX_VIDEO: regex::Regex = regex!(r"https?://(www\.|m\.)?(youtube\.[a-z]{2,3}/watch\?(feature=player_embedded&)?(list=[a-zA-Z0-9_-]+&)?v=[a-zA-Z0-9_-]+|youtu\.be/[a-zA-Z0-9_-]+)");
// https://regex101.com/r/aV1jS1/2
    pub static ref REGEX_PLAYLIST: regex::Regex = regex!(r"https?://(www\.|m\.)?youtube\.[a-z]{2,4}/(watch\?(feature=player_embedded&)?(v=[a-zA-Z0-9_-]+.*&)?list=[A-Za-z0-9_-]+|playlist\?list=[a-zA-Z0-9_-]+)");
}

/// Init youtube handler, registering it
pub fn init(registry: &mut Registry) {
    registry.register(Module {
        checker: Box::new(checker_file),
        handler: Box::new(handle_file),
    });
    registry.register(Module {
        checker: Box::new(checker_playlist),
        handler: Box::new(handle_playlist),
    });
}

/// Check if the data matches for the file handler
fn checker_file(data: &Request) -> bool {
    if data.playlist {
        return false;
    }
    REGEX_VIDEO.is_match(&data.url)
}

/// Check if the data matches for the playlist handler
fn checker_playlist(data: &Request) -> bool {
    if !data.playlist {
        return false;
    }
    REGEX_PLAYLIST.is_match(&data.url)
}

/// Playlist request handler for youtube
/// If compression is enabled all files will be downloaded into one dir and zipped afterwards
/// Otherwise for every entry in the playlist a new query is created. These will be handled one after another,
/// creating a file per entry.
fn handle_playlist(handle_db: &mut HandleData, request: &mut Request) -> Result<(), Error> {
    trace!("youtube playlist handler started");
    db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS);

    let name = Filename {
            name: try!(handle_db.downloader.get_playlist_name(&request.url)),
            extension: "zip".to_string(),
    };
    let mut step: i32 = 1;

    db::set_query_state(&mut request.get_conn(), &request.qid, "1/?");
    trace!("crawling ids");
    let file_ids = try!(handle_db.downloader.get_playlist_ids(request));

    if request.split {

        let save_path = try!(lib::format_save_path(&request.path, &name));;
        request.temp_path.push(request.qid.to_string());
        request.path = request.temp_path.clone();
        try!(create_dir(&request.path));
        let mut warnings = false;
        let mut current_url: String;
        let mut failed_log: String = String::from("Following urls couldn't be downloaded: \n");

        let max_steps = file_ids.len() as i32 + 2;
        db::update_steps(&mut request.get_conn(), &request.qid, 2,max_steps);
        for id in file_ids.iter() {
            step += 1;
            db::update_steps(&mut request.get_conn(), &request.qid, step,max_steps);
            current_url = String::from(YT_VIDEO_URL);
            current_url.push_str(&id);
            request.url = current_url.clone();
            match handle_file_int(handle_db, &request) {
                Err(e) => {
                    warn!("error downloading {}: {:?}", id, e);
                    failed_log.push_str(&format!("{} {:?}\n", current_url, e));
                    warnings = true;
                }
                Ok(_) => {}
            }
        }

        if warnings {
            debug!("found warnings");
            db::add_query_error(&mut request.get_conn(), &request.qid, &failed_log);
        }

        step += 1;
        db::update_steps(&mut request.get_conn(), &request.qid, step,max_steps);
        trace!("starting zipping");
        try!(lib::zip_folder(&request.temp_path, &save_path));
        trace!("adding file");
        handle_db.addFile(&save_path, &name.full_name());
        trace!("removing dir {}",request.path.to_string_lossy());
        try!(remove_dir_all(&request.path));
        trace!("updating state");

    } else {
        trace!("creating new requests for playlist entries");
        let mut current_url: String;
        for id in file_ids.iter() {
            current_url = String::from(YT_VIDEO_URL);
            current_url.push_str(&id);
            try!(db::add_sub_query(&current_url,&request));
        }
    }

    Ok(())
}

/// Handle file request
fn handle_file(handle_db: &mut HandleData, request: &mut Request) -> Result<(), Error> {
    try!(handle_file_int(handle_db, &request));

    Ok(())
}

/// Real handler, avoiding mut Request borrows for internal use
fn handle_file_int(handle_db: &mut HandleData, request: &Request) -> Result<(), Error> {
    trace!("youtube file handler started");
    if request.quality < 0 {
        try!(handle_audio(handle_db,request))
    } else {
        try!(handle_video(handle_db,request))
    }

    Ok(())
}

/// Handler for videos
fn handle_video(hdb: &mut HandleData, request: &Request) -> Result<(), Error> {
    condition!(db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS),!request.playlist);
    let mut temp_file_v = request.temp_path.clone();
    temp_file_v.push(request.qid.to_string());
    hdb.push(&temp_file_v);

    let mut dmca = false;

    trace!("Retriving name");
    let name = try!(get_name(&hdb, &request, true, true, &temp_file_v, &mut dmca));
    let origin_name = name.full_name();
    debug!("name: {}.{}", &name.name, &name.extension);

    let save_file = try!(lib::format_save_path(&request.path, &name));

    if dmca {
        try!(lib::move_file(&temp_file_v, &save_file));
        hdb.addFile(save_file.as_ref(), &origin_name);
        hdb.pop();
        return Ok(());
    }

    let audio_id = if name.extension == "mp4" {
        CONFIG.codecs.yt.audio_normal_mp4
    } else {
        CONFIG.codecs.yt.audio_normal_webm
    };

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 1,3),!request.playlist);
    trace!("downloading video");
    try!(hdb.downloader.download_file(&request, &temp_file_v, &request.quality.to_string()));

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 2,3),!request.playlist);
    let mut temp_file_a = request.temp_path.clone();
    temp_file_a.push(format!("{}a", request.qid));
    hdb.push(&temp_file_a);
    trace!("downloading audio");
    try!(hdb.downloader.download_file(&request, &temp_file_a, &audio_id.to_string()));

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 3,3),!request.playlist);
    trace!("merging");
    try!(hdb.converter.merge_files(&request.qid, &temp_file_v, &temp_file_a, &save_file,&mut request.get_conn()));
    if !request.playlist {
        hdb.addFile(&save_file, &origin_name);
    }
    try!(remove_file(&temp_file_a));
    hdb.pop();
    try!(remove_file(&temp_file_v));
    hdb.pop();

    Ok(())
}

/// Handler for audios
fn handle_audio(hdb: &mut HandleData, request: &Request) -> Result<(), Error> {
    condition!(db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS),!request.playlist);
    let mut dmca = false;
    let quality = try!(get_audio_quality(&request.quality));

    let mut temp_file_v = request.temp_path.clone();
    temp_file_v.push(request.qid.to_string());
    hdb.push(&temp_file_v);

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 1,3),!request.playlist);
    let mut name = try!(get_name(&hdb, &request, false, false, &temp_file_v, &mut dmca));

    if dmca {
        let save_file = try!(lib::format_save_path(&request.path, &name));
        try!(lib::move_file(&temp_file_v, &save_file));
        hdb.pop();
        hdb.addFile(&save_file, &name.full_name());
        return Ok(());
    }

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 2,3),!request.playlist);
    try!(hdb.downloader.download_file(&request, &temp_file_v, &quality));

    if request.quality == CONFIG.codecs.audio_raw ||
       request.quality == CONFIG.codecs.audio_source_hq {
        name.extension = String::from("m4a");
    } else if request.quality == CONFIG.codecs.audio_mp3 {
        name.extension = String::from("mp3");
    }

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 3,3),!request.playlist);
    let file = try!(lib::format_save_path(&request.path, &name));
    try!(hdb.converter.extract_audio(&request.qid,
                                     &temp_file_v,
                                     &file,
                                     request.quality == CONFIG.codecs.audio_mp3,
                                     &mut request.get_conn()));
    try!(remove_file(&temp_file_v));
    hdb.pop();
    if !request.playlist {
        hdb.addFile(&file, &name.full_name());
    }

    Ok(())
}

/// Retrive name of youtube video
/// If we should encounter an DMCA we'll let the lib call handle this
/// (if enabled), this gives us a downloaded file and the name
/// In this case we only need to move the file to it's destination
/// If use_quality is true, the name will be retrived for the request's quality
fn get_name<'a>(hdb: &HandleData,
                request: &Request,
                use_qality: bool,
                video: bool,
                path: &Path,
                dmca: &'a mut bool)
                -> Result<Filename, Error> {
    let quality = if use_qality {
        Some(request.quality.to_string())
    } else {
        None
    };
    Ok(match hdb.downloader.get_file_name(&request.url, quality) { // get filename
        Ok(v) => {
            v
        }
        Err(Error::DMCAError) => {
            // now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            info!("DMCA error!");
            if CONFIG.general.lib_use {
                if !request.split {
                    db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS);
                }
                match hdb.downloader.lib_request_video(1,
                                                       0,
                                                       path,
                                                       request,
                                                       &request.quality.to_string(),
                                                       video) {
                    Err(err) => {
                        warn!("lib-call error {:?}", err);
                        return Err(err);
                    }
                    Ok(v) => {
                        *dmca = true;
                        v
                    }
                }
            } else {
                return Err(Error::NotAvailable);
            }
        }
        Err(e) => {
            // unknown error / restricted source etc.. abort
            error!("Unknown error: {:?}", e);
            return Err(e);
        }
    })
}

/// Get audio quality to use
fn get_audio_quality(id: &i16) -> Result<String, Error> {
    let id = if *id == CONFIG.codecs.audio_raw {
        CONFIG.codecs.yt.audio_normal_mp4
    } else if *id == CONFIG.codecs.audio_source_hq {
        CONFIG.codecs.yt.audio_normal_mp4
    } else if *id == CONFIG.codecs.audio_mp3 {
        CONFIG.codecs.yt.audio_normal_mp4
    } else {
        return Err(Error::InputError("Unknown audio quality!".to_string()));
    };
    Ok(id.to_string())
}

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

    #[test]
    fn regex() {
        assert!(REGEX_VIDEO.is_match(r"https://m.youtube.com/watch?list=PLTXoSHLJey0RR60hjLhuUAaj_ftAdShqv&v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"http://m.youtube.com/watch?v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"https://m.youtube.com/watch?v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"http://youtu.be/IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA&list=PL6DA1502C5DDC0317&index=21"));
        assert!(!REGEX_VIDEO.is_match(r"https://www.youtube.com/playlist?list=PLJYiF4qyO-fpaYWfTylcFu3VGUCVK4xfz"));

        assert!(REGEX_PLAYLIST.is_match(r"https://m.youtube.com/playlist?list=PLTXoSHLJey0RR60hjLhuUAaj_ftAdShqv"));
        assert!(REGEX_PLAYLIST.is_match(r"https://m.youtube.com/watch?list=PLTXoSHLJey0RR60hjLhuUAaj_ftAdShqv&v=IO-_EoRSpUA"));
        assert!(REGEX_PLAYLIST.is_match(r"https://www.youtube.com/playlist?list=PLBCC15D0E3ED5E67A"));
        assert!(REGEX_PLAYLIST.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA&index=2&list=PLBCC15D0E3ED5E67A"));
        assert!(REGEX_PLAYLIST.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA&list=PL6DA1502C5DDC0317&index=21"));
        assert!(!REGEX_PLAYLIST.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA"));
    }

}
<commit_msg>Fix for playlist split handling & dmca playlist handling on youtube<commit_after>extern crate regex;

use super::{Registry, Module, HandleData};
use lib::{self, Error, Request, db};
use lib::downloader::Filename;
use std::fs::remove_file;
use std::fs::remove_dir_all;
use std::fs::create_dir;
use std::path::Path;

use CODE_IN_PROGRESS;
use CONFIG;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

macro_rules! condition(
    ($e:expr,$y:expr) => (match $y { true => {$e;}, false => ()} );
);

const YT_VIDEO_URL: &'static str = "https://www.youtube.com/watch?v=";

lazy_static! {
// https://regex101.com/r/lZ6lC1/3
// we need to remove the / escaping!
    pub static ref REGEX_VIDEO: regex::Regex = regex!(r"https?://(www\.|m\.)?(youtube\.[a-z]{2,3}/watch\?(feature=player_embedded&)?(list=[a-zA-Z0-9_-]+&)?v=[a-zA-Z0-9_-]+|youtu\.be/[a-zA-Z0-9_-]+)");
// https://regex101.com/r/aV1jS1/2
    pub static ref REGEX_PLAYLIST: regex::Regex = regex!(r"https?://(www\.|m\.)?youtube\.[a-z]{2,4}/(watch\?(feature=player_embedded&)?(v=[a-zA-Z0-9_-]+.*&)?list=[A-Za-z0-9_-]+|playlist\?list=[a-zA-Z0-9_-]+)");
}

/// Init youtube handler, registering it
pub fn init(registry: &mut Registry) {
    registry.register(Module {
        checker: Box::new(checker_file),
        handler: Box::new(handle_file),
    });
    registry.register(Module {
        checker: Box::new(checker_playlist),
        handler: Box::new(handle_playlist),
    });
}

/// Check if the data matches for the file handler
fn checker_file(data: &Request) -> bool {
    if data.playlist {
        return false;
    }
    REGEX_VIDEO.is_match(&data.url)
}

/// Check if the data matches for the playlist handler
fn checker_playlist(data: &Request) -> bool {
    if !data.playlist {
        return false;
    }
    REGEX_PLAYLIST.is_match(&data.url)
}

/// Playlist request handler for youtube
/// If compression is enabled all files will be downloaded into one dir and zipped afterwards
/// Otherwise for every entry in the playlist a new query is created. These will be handled one after another,
/// creating a file per entry.
fn handle_playlist(handle_db: &mut HandleData, request: &mut Request) -> Result<(), Error> {
    trace!("youtube playlist handler started");
    db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS);

    let name = Filename {
            name: try!(handle_db.downloader.get_playlist_name(&request.url)),
            extension: "zip".to_string(),
    };
    let mut step: i32 = 1;

    db::set_query_state(&mut request.get_conn(), &request.qid, "1/?");
    trace!("crawling ids");
    let file_ids = try!(handle_db.downloader.get_playlist_ids(request));

    if request.split {
        trace!("creating new requests for playlist entries");
        let mut current_url: String;
        for id in file_ids.iter() {
            current_url = String::from(YT_VIDEO_URL);
            current_url.push_str(&id);
            try!(db::add_sub_query(&current_url,&request));
        }
    } else {
        let save_path = try!(lib::format_save_path(&request.path, &name));;
        request.temp_path.push(request.qid.to_string());
        request.path = request.temp_path.clone();
        try!(create_dir(&request.path));
        let mut warnings = false;
        let mut current_url: String;
        let mut failed_log: String = String::from("Following urls couldn't be downloaded: \n");

        let max_steps = file_ids.len() as i32 + 2;
        db::update_steps(&mut request.get_conn(), &request.qid, 2,max_steps);
        for id in file_ids.iter() {
            step += 1;
            db::update_steps(&mut request.get_conn(), &request.qid, step,max_steps);
            current_url = String::from(YT_VIDEO_URL);
            current_url.push_str(&id);
            request.url = current_url.clone();
            match handle_file_int(handle_db, &request) {
                Err(e) => {
                    warn!("error downloading {}: {:?}", id, e);
                    failed_log.push_str(&format!("{} {:?}\n", current_url, e));
                    warnings = true;
                }
                Ok(_) => {}
            }
        }

        if warnings {
            debug!("found warnings");
            db::add_query_error(&mut request.get_conn(), &request.qid, &failed_log);
        }

        step += 1;
        db::update_steps(&mut request.get_conn(), &request.qid, step,max_steps);
        trace!("starting zipping");
        try!(lib::zip_folder(&request.temp_path, &save_path));
        trace!("adding file");
        handle_db.addFile(&save_path, &name.full_name());
        trace!("removing dir {}",request.path.to_string_lossy());
        try!(remove_dir_all(&request.path));
        trace!("updating state");
    }

    Ok(())
}

/// Handle file request
fn handle_file(handle_db: &mut HandleData, request: &mut Request) -> Result<(), Error> {
    try!(handle_file_int(handle_db, &request));

    Ok(())
}

/// Real handler, avoiding mut Request borrows for internal use
fn handle_file_int(handle_db: &mut HandleData, request: &Request) -> Result<(), Error> {
    trace!("youtube file handler started");
    if request.quality < 0 {
        try!(handle_audio(handle_db,request))
    } else {
        try!(handle_video(handle_db,request))
    }

    Ok(())
}

/// Handler for videos
fn handle_video(hdb: &mut HandleData, request: &Request) -> Result<(), Error> {
    condition!(db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS),!request.playlist);
    let mut temp_file_v = request.temp_path.clone();
    temp_file_v.push(request.qid.to_string());
    hdb.push(&temp_file_v);

    let mut dmca = false;

    trace!("Retriving name");
    let name = try!(get_name(&hdb, &request, true, true, &temp_file_v, &mut dmca));
    let origin_name = name.full_name();
    debug!("name: {}.{}", &name.name, &name.extension);

    let save_file = try!(lib::format_save_path(&request.path, &name));

    if dmca {
        try!(lib::move_file(&temp_file_v, &save_file));
        hdb.addFile(save_file.as_ref(), &origin_name);
        hdb.pop();
        return Ok(());
    }

    let audio_id = if name.extension == "mp4" {
        CONFIG.codecs.yt.audio_normal_mp4
    } else {
        CONFIG.codecs.yt.audio_normal_webm
    };

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 1,3),!request.playlist);
    trace!("downloading video");
    try!(hdb.downloader.download_file(&request, &temp_file_v, &request.quality.to_string()));

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 2,3),!request.playlist);
    let mut temp_file_a = request.temp_path.clone();
    temp_file_a.push(format!("{}a", request.qid));
    hdb.push(&temp_file_a);
    trace!("downloading audio");
    try!(hdb.downloader.download_file(&request, &temp_file_a, &audio_id.to_string()));

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 3,3),!request.playlist);
    trace!("merging");
    try!(hdb.converter.merge_files(&request.qid, &temp_file_v, &temp_file_a, &save_file,&mut request.get_conn()));
    if !request.playlist {
        hdb.addFile(&save_file, &origin_name);
    }
    try!(remove_file(&temp_file_a));
    hdb.pop();
    try!(remove_file(&temp_file_v));
    hdb.pop();

    Ok(())
}

/// Handler for audios
fn handle_audio(hdb: &mut HandleData, request: &Request) -> Result<(), Error> {
    condition!(db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS),!request.playlist);
    let mut dmca = false;
    let quality = try!(get_audio_quality(&request.quality));

    let mut temp_file_v = request.temp_path.clone();
    temp_file_v.push(request.qid.to_string());
    hdb.push(&temp_file_v);

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 1,3),!request.playlist);
    let mut name = try!(get_name(&hdb, &request, false, false, &temp_file_v, &mut dmca));

    if dmca {
        let save_file = try!(lib::format_save_path(&request.path, &name));
        try!(lib::move_file(&temp_file_v, &save_file));
        if !request.playlist {
            hdb.pop();
            hdb.addFile(&save_file, &name.full_name());
        }
        return Ok(());
    }

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 2,3),!request.playlist);
    try!(hdb.downloader.download_file(&request, &temp_file_v, &quality));

    if request.quality == CONFIG.codecs.audio_raw ||
       request.quality == CONFIG.codecs.audio_source_hq {
        name.extension = String::from("m4a");
    } else if request.quality == CONFIG.codecs.audio_mp3 {
        name.extension = String::from("mp3");
    }

    condition!(db::update_steps(&mut request.get_conn(), &request.qid, 3,3),!request.playlist);
    let file = try!(lib::format_save_path(&request.path, &name));
    try!(hdb.converter.extract_audio(&request.qid,
                                     &temp_file_v,
                                     &file,
                                     request.quality == CONFIG.codecs.audio_mp3,
                                     &mut request.get_conn()));
    try!(remove_file(&temp_file_v));
    hdb.pop();
    if !request.playlist {
        hdb.addFile(&file, &name.full_name());
    }

    Ok(())
}

/// Retrive name of youtube video
/// If we should encounter an DMCA we'll let the lib call handle this
/// (if enabled), this gives us a downloaded file and the name
/// In this case we only need to move the file to it's destination
/// If use_quality is true, the name will be retrived for the request's quality
fn get_name<'a>(hdb: &HandleData,
                request: &Request,
                use_qality: bool,
                video: bool,
                path: &Path,
                dmca: &'a mut bool)
                -> Result<Filename, Error> {
    let quality = if use_qality {
        Some(request.quality.to_string())
    } else {
        None
    };
    Ok(match hdb.downloader.get_file_name(&request.url, quality) { // get filename
        Ok(v) => {
            v
        }
        Err(Error::DMCAError) => {
            // now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            info!("DMCA error!");
            if CONFIG.general.lib_use {
                if !request.split {
                    db::set_query_code(&mut request.get_conn(), &request.qid, &CODE_IN_PROGRESS);
                }
                match hdb.downloader.lib_request_video(1,
                                                       0,
                                                       path,
                                                       request,
                                                       &request.quality.to_string(),
                                                       video) {
                    Err(err) => {
                        warn!("lib-call error {:?}", err);
                        return Err(err);
                    }
                    Ok(v) => {
                        *dmca = true;
                        v
                    }
                }
            } else {
                return Err(Error::NotAvailable);
            }
        }
        Err(e) => {
            // unknown error / restricted source etc.. abort
            error!("Unknown error: {:?}", e);
            return Err(e);
        }
    })
}

/// Get audio quality to use
fn get_audio_quality(id: &i16) -> Result<String, Error> {
    let id = if *id == CONFIG.codecs.audio_raw {
        CONFIG.codecs.yt.audio_normal_mp4
    } else if *id == CONFIG.codecs.audio_source_hq {
        CONFIG.codecs.yt.audio_normal_mp4
    } else if *id == CONFIG.codecs.audio_mp3 {
        CONFIG.codecs.yt.audio_normal_mp4
    } else {
        return Err(Error::InputError("Unknown audio quality!".to_string()));
    };
    Ok(id.to_string())
}

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

    #[test]
    fn regex() {
        assert!(REGEX_VIDEO.is_match(r"https://m.youtube.com/watch?list=PLTXoSHLJey0RR60hjLhuUAaj_ftAdShqv&v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"http://m.youtube.com/watch?v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"https://m.youtube.com/watch?v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"http://youtu.be/IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA"));
        assert!(REGEX_VIDEO.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA&list=PL6DA1502C5DDC0317&index=21"));
        assert!(!REGEX_VIDEO.is_match(r"https://www.youtube.com/playlist?list=PLJYiF4qyO-fpaYWfTylcFu3VGUCVK4xfz"));

        assert!(REGEX_PLAYLIST.is_match(r"https://m.youtube.com/playlist?list=PLTXoSHLJey0RR60hjLhuUAaj_ftAdShqv"));
        assert!(REGEX_PLAYLIST.is_match(r"https://m.youtube.com/watch?list=PLTXoSHLJey0RR60hjLhuUAaj_ftAdShqv&v=IO-_EoRSpUA"));
        assert!(REGEX_PLAYLIST.is_match(r"https://www.youtube.com/playlist?list=PLBCC15D0E3ED5E67A"));
        assert!(REGEX_PLAYLIST.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA&index=2&list=PLBCC15D0E3ED5E67A"));
        assert!(REGEX_PLAYLIST.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA&list=PL6DA1502C5DDC0317&index=21"));
        assert!(!REGEX_PLAYLIST.is_match(r"https://www.youtube.com/watch?v=IO-_EoRSpUA"));
    }

}
<|endoftext|>


<commit_before>use crate::ptrace_control::*;
use crate::statemachine::*;
use nix::unistd::Pid;
use nix::{Error, Result};
use std::collections::HashMap;

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
const INT: u64 = 0xCC;

/// Breakpoint construct used to monitor program execution. As tarpaulin is an
/// automated process, this will likely have less functionality than most
/// breakpoint implementations.
#[derive(Debug)]
pub struct Breakpoint {
    /// Program counter
    pub pc: u64,
    /// Bottom byte of address data.
    /// This is replaced to enable the interrupt. Rest of data is never changed.
    data: u8,
    /// Reading from memory with ptrace gives addresses aligned to bytes.
    /// We therefore need to know the shift to place the breakpoint in the right place
    shift: u64,
    /// Map of the state of the breakpoint on each thread/process
    is_running: HashMap<Pid, bool>,
}

impl Breakpoint {
    /// Creates a new breakpoint for the given process and program counter.
    pub fn new(pid: Pid, pc: u64) -> Result<Self> {
        let aligned = align_address(pc);
        let data = read_address(pid, aligned)?;
        let shift = 8 * (pc - aligned);
        let data = ((data >> shift) & 0xFF) as u8;

        let mut b = Breakpoint {
            pc,
            data,
            shift,
            is_running: HashMap::new(),
        };
        match b.enable(pid) {
            Ok(_) => Ok(b),
            Err(e) => Err(e),
        }
    }

    pub fn jump_to(&mut self, pid: Pid) -> Result<()> {
        set_instruction_pointer(pid, self.pc).map(|_| ())
    }

    /// Attaches the current breakpoint.
    pub fn enable(&mut self, pid: Pid) -> Result<()> {
        let data = read_address(pid, self.aligned_address())?;
        self.is_running.insert(pid, true);
        let mut intdata = data & (!(0xFFu64 << self.shift) as i64);
        intdata |= (INT << self.shift) as i64;
        if data == intdata {
            Err(Error::UnknownErrno)
        } else {
            write_to_address(pid, self.aligned_address(), intdata)
        }
    }

    pub fn disable(&self, pid: Pid) -> Result<()> {
        // I require the bit fiddlin this end.
        let data = read_address(pid, self.aligned_address())?;
        let mut orgdata = data & (!(0xFFu64 << self.shift) as i64);
        orgdata |= i64::from(self.data) << self.shift;
        write_to_address(pid, self.aligned_address(), orgdata)
    }

    /// Processes the breakpoint. This steps over the breakpoint
    pub fn process(
        &mut self,
        pid: Pid,
        reenable: bool,
    ) -> Result<(bool, TracerAction<ProcessInfo>)> {
        let is_running = match self.is_running.get(&pid) {
            Some(r) => *r,
            None => true,
        };
        if is_running {
            let _ = self.enable(pid);
            self.step(pid)?;
            self.is_running.insert(pid, false);
            Ok((true, TracerAction::Step(pid.into())))
        } else {
            self.disable(pid)?;
            if reenable {
                self.enable(pid)?;
            }
            self.is_running.insert(pid, true);
            Ok((false, TracerAction::Continue(pid.into())))
        }
    }

    /// Call this when a ptrace thread is killed. Won't reenable the breakpoint
    /// so may lose the ability to instrument this line.
    pub fn thread_killed(&mut self, pid: Pid) {
        self.is_running.remove(&pid);
    }

    /// Steps past the current breakpoint.
    /// For more advanced coverage may interrogate the variables of a branch.
    fn step(&mut self, pid: Pid) -> Result<()> {
        // Remove the breakpoint, reset the program counter to step before it
        // hit the breakpoint then step to execute the original instruction.
        self.disable(pid)?;
        self.jump_to(pid)?;
        Ok(())
    }

    pub fn aligned_address(&self) -> u64 {
        align_address(self.pc)
    }
}

#[inline(always)]
pub(crate) fn align_address(addr: u64) -> u64 {
    addr & !0x7u64
}
<commit_msg>Update comment on INT<commit_after>use crate::ptrace_control::*;
use crate::statemachine::*;
use nix::unistd::Pid;
use nix::{Error, Result};
use std::collections::HashMap;

/// INT refers to the software interrupt instruction. For x64/x86 we use INT3 which is a
/// one byte instruction defined for use by debuggers. For implementing support for other
/// architectures the equivalent instructions will have to be found and also the architectures
/// added to the CI.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
const INT: u64 = 0xCC;

/// Breakpoint construct used to monitor program execution. As tarpaulin is an
/// automated process, this will likely have less functionality than most
/// breakpoint implementations.
#[derive(Debug)]
pub struct Breakpoint {
    /// Program counter
    pub pc: u64,
    /// Bottom byte of address data.
    /// This is replaced to enable the interrupt. Rest of data is never changed.
    data: u8,
    /// Reading from memory with ptrace gives addresses aligned to bytes.
    /// We therefore need to know the shift to place the breakpoint in the right place
    shift: u64,
    /// Map of the state of the breakpoint on each thread/process
    is_running: HashMap<Pid, bool>,
}

impl Breakpoint {
    /// Creates a new breakpoint for the given process and program counter.
    pub fn new(pid: Pid, pc: u64) -> Result<Self> {
        let aligned = align_address(pc);
        let data = read_address(pid, aligned)?;
        let shift = 8 * (pc - aligned);
        let data = ((data >> shift) & 0xFF) as u8;

        let mut b = Breakpoint {
            pc,
            data,
            shift,
            is_running: HashMap::new(),
        };
        match b.enable(pid) {
            Ok(_) => Ok(b),
            Err(e) => Err(e),
        }
    }

    pub fn jump_to(&mut self, pid: Pid) -> Result<()> {
        set_instruction_pointer(pid, self.pc).map(|_| ())
    }

    /// Attaches the current breakpoint.
    pub fn enable(&mut self, pid: Pid) -> Result<()> {
        let data = read_address(pid, self.aligned_address())?;
        self.is_running.insert(pid, true);
        let mut intdata = data & (!(0xFFu64 << self.shift) as i64);
        intdata |= (INT << self.shift) as i64;
        if data == intdata {
            Err(Error::UnknownErrno)
        } else {
            write_to_address(pid, self.aligned_address(), intdata)
        }
    }

    pub fn disable(&self, pid: Pid) -> Result<()> {
        // I require the bit fiddlin this end.
        let data = read_address(pid, self.aligned_address())?;
        let mut orgdata = data & (!(0xFFu64 << self.shift) as i64);
        orgdata |= i64::from(self.data) << self.shift;
        write_to_address(pid, self.aligned_address(), orgdata)
    }

    /// Processes the breakpoint. This steps over the breakpoint
    pub fn process(
        &mut self,
        pid: Pid,
        reenable: bool,
    ) -> Result<(bool, TracerAction<ProcessInfo>)> {
        let is_running = match self.is_running.get(&pid) {
            Some(r) => *r,
            None => true,
        };
        if is_running {
            let _ = self.enable(pid);
            self.step(pid)?;
            self.is_running.insert(pid, false);
            Ok((true, TracerAction::Step(pid.into())))
        } else {
            self.disable(pid)?;
            if reenable {
                self.enable(pid)?;
            }
            self.is_running.insert(pid, true);
            Ok((false, TracerAction::Continue(pid.into())))
        }
    }

    /// Call this when a ptrace thread is killed. Won't reenable the breakpoint
    /// so may lose the ability to instrument this line.
    pub fn thread_killed(&mut self, pid: Pid) {
        self.is_running.remove(&pid);
    }

    /// Steps past the current breakpoint.
    /// For more advanced coverage may interrogate the variables of a branch.
    fn step(&mut self, pid: Pid) -> Result<()> {
        // Remove the breakpoint, reset the program counter to step before it
        // hit the breakpoint then step to execute the original instruction.
        self.disable(pid)?;
        self.jump_to(pid)?;
        Ok(())
    }

    pub fn aligned_address(&self) -> u64 {
        align_address(self.pc)
    }
}

#[inline(always)]
pub(crate) fn align_address(addr: u64) -> u64 {
    addr & !0x7u64
}
<|endoftext|>


<commit_before>/// Return early with an error.
///
/// This macro is equivalent to `return Err(`[`anyhow!($args...)`][anyhow!]`)`.
///
/// The surrounding function's or closure's return value is required to be
/// `Result<_,`[`anyhow::Error`][crate::Error]`>`.
///
/// # Example
///
/// ```
/// # use anyhow::{bail, Result};
/// #
/// # fn has_permission(user: usize, resource: usize) -> bool {
/// #     true
/// # }
/// #
/// # fn main() -> Result<()> {
/// #     let user = 0;
/// #     let resource = 0;
/// #
/// if !has_permission(user, resource) {
///     bail!("permission denied for accessing {}", resource);
/// }
/// #     Ok(())
/// # }
/// ```
///
/// ```
/// # use anyhow::{bail, Result};
/// # use thiserror::Error;
/// #
/// # const MAX_DEPTH: usize = 1;
/// #
/// #[derive(Error, Debug)]
/// enum ScienceError {
///     #[error("recursion limit exceeded")]
///     RecursionLimitExceeded,
///     # #[error("...")]
///     # More = (stringify! {
///     ...
///     # }, 1).1,
/// }
///
/// # fn main() -> Result<()> {
/// #     let depth = 0;
/// #
/// if depth > MAX_DEPTH {
///     bail!(ScienceError::RecursionLimitExceeded);
/// }
/// #     Ok(())
/// # }
/// ```
#[cfg(doc)]
#[macro_export]
macro_rules! bail {
    ($msg:literal $(,)?) => {
        return $crate::private::Err($crate::__anyhow!($msg))
    };
    ($err:expr $(,)?) => {
        return $crate::private::Err($crate::__anyhow!($err))
    };
    ($fmt:expr, $($arg:tt)*) => {
        return $crate::private::Err($crate::__anyhow!($fmt, $($arg)*))
    };
}

// Workaround for crates that intentionally contained `{}` in an error message
// prior to https://github.com/dtolnay/anyhow/issues/55 catching the missing
// format args.
#[cfg(not(doc))]
#[macro_export]
macro_rules! bail {
    // https://github.com/estk/log4rs/blob/afa0351af56b3bfd1780389700051d7e4d8bbdc9/src/append/rolling_file/policy/compound/roll/fixed_window.rs#L261
    ("pattern does not contain `{}`") => {
        return $crate::private::Err($crate::Error::msg("pattern does not contain `{}`"))
    };
    ($msg:literal $(,)?) => {
        return $crate::private::Err($crate::__anyhow!($msg))
    };
    ($err:expr $(,)?) => {
        return $crate::private::Err($crate::__anyhow!($err))
    };
    ($fmt:expr, $($arg:tt)*) => {
        return $crate::private::Err($crate::__anyhow!($fmt, $($arg)*))
    };
}

/// Return early with an error if a condition is not satisfied.
///
/// This macro is equivalent to `if !$cond { return
/// Err(`[`anyhow!($args...)`][anyhow!]`); }`.
///
/// The surrounding function's or closure's return value is required to be
/// `Result<_,`[`anyhow::Error`][crate::Error]`>`.
///
/// Analogously to `assert!`, `ensure!` takes a condition and exits the function
/// if the condition fails. Unlike `assert!`, `ensure!` returns an `Error`
/// rather than panicking.
///
/// # Example
///
/// ```
/// # use anyhow::{ensure, Result};
/// #
/// # fn main() -> Result<()> {
/// #     let user = 0;
/// #
/// ensure!(user == 0, "only user 0 is allowed");
/// #     Ok(())
/// # }
/// ```
///
/// ```
/// # use anyhow::{ensure, Result};
/// # use thiserror::Error;
/// #
/// # const MAX_DEPTH: usize = 1;
/// #
/// #[derive(Error, Debug)]
/// enum ScienceError {
///     #[error("recursion limit exceeded")]
///     RecursionLimitExceeded,
///     # #[error("...")]
///     # More = (stringify! {
///     ...
///     # }, 1).1,
/// }
///
/// # fn main() -> Result<()> {
/// #     let depth = 0;
/// #
/// ensure!(depth <= MAX_DEPTH, ScienceError::RecursionLimitExceeded);
/// #     Ok(())
/// # }
/// ```
#[cfg(doc)]
#[macro_export]
macro_rules! ensure {
    ($cond:expr $(,)?) => {
        if !$cond {
            return $crate::private::Err($crate::Error::msg(
                $crate::private::concat!("Condition failed: `", $crate::private::stringify!($cond), "`")
            ));
        }
    };
    ($cond:expr, $msg:literal $(,)?) => {
        if !$cond {
            return $crate::private::Err($crate::__anyhow!($msg));
        }
    };
    ($cond:expr, $err:expr $(,)?) => {
        if !$cond {
            return $crate::private::Err($crate::__anyhow!($err));
        }
    };
    ($cond:expr, $fmt:expr, $($arg:tt)*) => {
        if !$cond {
            return $crate::private::Err($crate::__anyhow!($fmt, $($arg)*));
        }
    };
}

#[cfg(not(doc))]
#[macro_export]
macro_rules! ensure {
    ($($tt:tt)*) => {
        $crate::__parse_ensure!(
            /* state */ 0
            /* stack */ ()
            /* bail */ ($($tt)*)
            /* fuel */ (~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~)
            /* parse */ {()}
            /* dup */ ($($tt)*)
            /* rest */ $($tt)*
        )
    };
}

/// Construct an ad-hoc error from a string or existing non-`anyhow` error
/// value.
///
/// This evaluates to an [`Error`][crate::Error]. It can take either just a
/// string, or a format string with arguments. It also can take any custom type
/// which implements `Debug` and `Display`.
///
/// If called with a single argument whose type implements `std::error::Error`
/// (in addition to `Debug` and `Display`, which are always required), then that
/// Error impl's `source` is preserved as the `source` of the resulting
/// `anyhow::Error`.
///
/// # Example
///
/// ```
/// # type V = ();
/// #
/// use anyhow::{anyhow, Result};
///
/// fn lookup(key: &str) -> Result<V> {
///     if key.len() != 16 {
///         return Err(anyhow!("key length must be 16 characters, got {:?}", key));
///     }
///
///     // ...
///     # Ok(())
/// }
/// ```
#[macro_export]
macro_rules! anyhow {
    ($msg:literal $(,)?) => {
        $crate::private::must_use({
            let error = $crate::private::format_err($crate::private::format_args!($msg));
            error
        })
    };
    ($err:expr $(,)?) => {
        $crate::private::must_use({
            use $crate::private::kind::*;
            let error = match $err {
                error => (&error).anyhow_kind().new(error),
            };
            error
        })
    };
    ($fmt:expr, $($arg:tt)*) => {
        $crate::Error::msg($crate::private::format!($fmt, $($arg)*))
    };
}

// Not public API. This is used in the implementation of some of the other
// macros, in which the must_use call is not needed because the value is known
// to be used.
#[doc(hidden)]
#[macro_export]
macro_rules! __anyhow {
    ($msg:literal $(,)?) => ({
        let error = $crate::private::format_err($crate::private::format_args!($msg));
        error
    });
    ($err:expr $(,)?) => ({
        use $crate::private::kind::*;
        let error = match $err {
            error => (&error).anyhow_kind().new(error),
        };
        error
    });
    ($fmt:expr, $($arg:tt)*) => {
        $crate::Error::msg($crate::private::format!($fmt, $($arg)*))
    };
}
<commit_msg>Remove log4rs workaround from bail macro<commit_after>/// Return early with an error.
///
/// This macro is equivalent to `return Err(`[`anyhow!($args...)`][anyhow!]`)`.
///
/// The surrounding function's or closure's return value is required to be
/// `Result<_,`[`anyhow::Error`][crate::Error]`>`.
///
/// # Example
///
/// ```
/// # use anyhow::{bail, Result};
/// #
/// # fn has_permission(user: usize, resource: usize) -> bool {
/// #     true
/// # }
/// #
/// # fn main() -> Result<()> {
/// #     let user = 0;
/// #     let resource = 0;
/// #
/// if !has_permission(user, resource) {
///     bail!("permission denied for accessing {}", resource);
/// }
/// #     Ok(())
/// # }
/// ```
///
/// ```
/// # use anyhow::{bail, Result};
/// # use thiserror::Error;
/// #
/// # const MAX_DEPTH: usize = 1;
/// #
/// #[derive(Error, Debug)]
/// enum ScienceError {
///     #[error("recursion limit exceeded")]
///     RecursionLimitExceeded,
///     # #[error("...")]
///     # More = (stringify! {
///     ...
///     # }, 1).1,
/// }
///
/// # fn main() -> Result<()> {
/// #     let depth = 0;
/// #
/// if depth > MAX_DEPTH {
///     bail!(ScienceError::RecursionLimitExceeded);
/// }
/// #     Ok(())
/// # }
/// ```
#[macro_export]
macro_rules! bail {
    ($msg:literal $(,)?) => {
        return $crate::private::Err($crate::__anyhow!($msg))
    };
    ($err:expr $(,)?) => {
        return $crate::private::Err($crate::__anyhow!($err))
    };
    ($fmt:expr, $($arg:tt)*) => {
        return $crate::private::Err($crate::__anyhow!($fmt, $($arg)*))
    };
}

/// Return early with an error if a condition is not satisfied.
///
/// This macro is equivalent to `if !$cond { return
/// Err(`[`anyhow!($args...)`][anyhow!]`); }`.
///
/// The surrounding function's or closure's return value is required to be
/// `Result<_,`[`anyhow::Error`][crate::Error]`>`.
///
/// Analogously to `assert!`, `ensure!` takes a condition and exits the function
/// if the condition fails. Unlike `assert!`, `ensure!` returns an `Error`
/// rather than panicking.
///
/// # Example
///
/// ```
/// # use anyhow::{ensure, Result};
/// #
/// # fn main() -> Result<()> {
/// #     let user = 0;
/// #
/// ensure!(user == 0, "only user 0 is allowed");
/// #     Ok(())
/// # }
/// ```
///
/// ```
/// # use anyhow::{ensure, Result};
/// # use thiserror::Error;
/// #
/// # const MAX_DEPTH: usize = 1;
/// #
/// #[derive(Error, Debug)]
/// enum ScienceError {
///     #[error("recursion limit exceeded")]
///     RecursionLimitExceeded,
///     # #[error("...")]
///     # More = (stringify! {
///     ...
///     # }, 1).1,
/// }
///
/// # fn main() -> Result<()> {
/// #     let depth = 0;
/// #
/// ensure!(depth <= MAX_DEPTH, ScienceError::RecursionLimitExceeded);
/// #     Ok(())
/// # }
/// ```
#[cfg(doc)]
#[macro_export]
macro_rules! ensure {
    ($cond:expr $(,)?) => {
        if !$cond {
            return $crate::private::Err($crate::Error::msg(
                $crate::private::concat!("Condition failed: `", $crate::private::stringify!($cond), "`")
            ));
        }
    };
    ($cond:expr, $msg:literal $(,)?) => {
        if !$cond {
            return $crate::private::Err($crate::__anyhow!($msg));
        }
    };
    ($cond:expr, $err:expr $(,)?) => {
        if !$cond {
            return $crate::private::Err($crate::__anyhow!($err));
        }
    };
    ($cond:expr, $fmt:expr, $($arg:tt)*) => {
        if !$cond {
            return $crate::private::Err($crate::__anyhow!($fmt, $($arg)*));
        }
    };
}

#[cfg(not(doc))]
#[macro_export]
macro_rules! ensure {
    ($($tt:tt)*) => {
        $crate::__parse_ensure!(
            /* state */ 0
            /* stack */ ()
            /* bail */ ($($tt)*)
            /* fuel */ (~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~ ~~~~~~~~~~)
            /* parse */ {()}
            /* dup */ ($($tt)*)
            /* rest */ $($tt)*
        )
    };
}

/// Construct an ad-hoc error from a string or existing non-`anyhow` error
/// value.
///
/// This evaluates to an [`Error`][crate::Error]. It can take either just a
/// string, or a format string with arguments. It also can take any custom type
/// which implements `Debug` and `Display`.
///
/// If called with a single argument whose type implements `std::error::Error`
/// (in addition to `Debug` and `Display`, which are always required), then that
/// Error impl's `source` is preserved as the `source` of the resulting
/// `anyhow::Error`.
///
/// # Example
///
/// ```
/// # type V = ();
/// #
/// use anyhow::{anyhow, Result};
///
/// fn lookup(key: &str) -> Result<V> {
///     if key.len() != 16 {
///         return Err(anyhow!("key length must be 16 characters, got {:?}", key));
///     }
///
///     // ...
///     # Ok(())
/// }
/// ```
#[macro_export]
macro_rules! anyhow {
    ($msg:literal $(,)?) => {
        $crate::private::must_use({
            let error = $crate::private::format_err($crate::private::format_args!($msg));
            error
        })
    };
    ($err:expr $(,)?) => {
        $crate::private::must_use({
            use $crate::private::kind::*;
            let error = match $err {
                error => (&error).anyhow_kind().new(error),
            };
            error
        })
    };
    ($fmt:expr, $($arg:tt)*) => {
        $crate::Error::msg($crate::private::format!($fmt, $($arg)*))
    };
}

// Not public API. This is used in the implementation of some of the other
// macros, in which the must_use call is not needed because the value is known
// to be used.
#[doc(hidden)]
#[macro_export]
macro_rules! __anyhow {
    ($msg:literal $(,)?) => ({
        let error = $crate::private::format_err($crate::private::format_args!($msg));
        error
    });
    ($err:expr $(,)?) => ({
        use $crate::private::kind::*;
        let error = match $err {
            error => (&error).anyhow_kind().new(error),
        };
        error
    });
    ($fmt:expr, $($arg:tt)*) => {
        $crate::Error::msg($crate::private::format!($fmt, $($arg)*))
    };
}
<|endoftext|>


<commit_before>use std::env;
use std::ffi::CString;
use std::fmt;
use std::marker::PhantomData;
use std::mem;
use std::io::{Error, ErrorKind};
use std::io::Result as IoResult;
use std::iter::{Iterator, IntoIterator};
use std::path::Path;
use std::process::Command;

use cbox::{CBox, CSemiBox};
use ffi::analysis::LLVMVerifierFailureAction;
use ffi::{analysis, core, linker, LLVMModule};
use ffi::bit_writer as writer;
use ffi::bit_reader as reader;
use ffi::ir_reader as ir_reader;
use ffi::prelude::{LLVMValueRef, LLVMModuleRef};
use ffi::transforms::pass_manager_builder as builder;
use libc::{c_char, c_uint};

use buffer::MemoryBuffer;
use context::{Context, GetContext};
use util;
use ty::Type;
use value::{Function, GlobalValue, Value};

/// Represents a single compilation unit of code.
///
/// This is attached to the lifetime of the context that constructs it, but is owned by the `CSemiBox`.
pub struct Module;
native_ref!(&Module = LLVMModuleRef);
impl Module
{
  /// Create a new module in the context given with the name given.
  ///
  /// The lifetime of the module will match the lifetime of the context
  /// you instance it in because the context contains it.
  ///
  /// ```rust
  /// use llvm::*;
  /// let context = Context::new();
  /// let context = context.as_semi();
  /// let module = Module::new("name", &context);
  /// println!("{:?}", module)
  /// ```
  pub fn new<'a>(name: &str, context: &'a Context) -> CSemiBox<'a, Module>
  {
    let c_name = CString::new(name).unwrap();
    unsafe {
    	CSemiBox::new(core::LLVMModuleCreateWithNameInContext(c_name.as_ptr(), context.into()))
    }
  }

  /// Add an external global to the module with the given type and name.
  pub fn add_global<'a>(&'a self, name: &str, ty: &'a Type) -> &'a GlobalValue {
    util::with_cstr(name, |ptr| unsafe {
        core::LLVMAddGlobal(self.into(), ty.into(), ptr).into()
    })
  }

  /// Add a global in the given address space to the module with the given type and name.
  pub fn add_global_in_addr_space<'a>(&'a self,
  	                                  name: &str,
  	                                  ty: &'a Type,
  	                                  sp: AddressSpace) -> &'a GlobalValue
  {
    util::with_cstr(name, |ptr| unsafe {
      core::LLVMAddGlobalInAddressSpace(self.into(), ty.into(), ptr, sp as c_uint).into()
    })
  }

  /// Add a constant global to the module with the given type, name and value.
  pub fn add_global_constant<'a>(&'a self, name: &str, val: &'a Value) -> &'a GlobalValue
  {
    util::with_cstr(name, |ptr| unsafe {
      let global = core::LLVMAddGlobal(self.into(), val.get_type().into(), ptr);
      core::LLVMSetInitializer (global.into(), val.into());
      global.into()
    })
  }
  /// Get the global with the name given, or `None` if no global with that name exists.
  pub fn get_global<'a>(&'a self, name: &str) -> Option<&'a GlobalValue>
  {
    util::with_cstr(name, |ptr| unsafe {
        let ptr = core::LLVMGetNamedGlobal(self.into(), ptr);
        util::ptr_to_null(ptr)
    })
  }

  pub fn parse_ir<'a>(context: &'a Context, path: &str) -> Result<CSemiBox<'a, Module>, CBox<str>>
  {
  	unsafe {
      let mut out = mem::uninitialized();
      let mut err = mem::uninitialized();
      let buf = try!(MemoryBuffer::new_from_file(path));

      let ret = ir_reader::LLVMParseIRInContext(context.into(), buf.as_ptr(), &mut out, &mut err);
      if ret == 1 {
      	Err(CBox::new(err))
      } else {
        Ok(CSemiBox::new(out))
      }
    }
  }

  /// Parse this bitcode file into a module, or return an error string.
  pub fn parse_bitcode<'a>(context: &'a Context, path: &str) -> Result<CSemiBox<'a, Module>, CBox<str>> {
    unsafe {
      let mut out = mem::uninitialized();
      let mut err = mem::uninitialized();
      let buf = try!(MemoryBuffer::new_from_file(path));

      let ret = reader::LLVMParseBitcodeInContext(context.into(),
      																					  buf.as_ptr(),
      	                                          &mut out,
      	                                          &mut err);
      if ret == 1 {
          Err(CBox::new(err))
      } else {
          Ok(CSemiBox::new(out))
      }
    }
  }
  /// Write this module's bitcode to the path given.
  pub fn write_bitcode(&self, path: &str) -> IoResult<()>
  {
     util::with_cstr(path, |cpath| unsafe {
       if writer::LLVMWriteBitcodeToFile(self.into(), cpath) != 0 {
           Err(Error::new(ErrorKind::Other, &format!("could not write to {}", path) as &str))
       } else {
           Ok(())
       }
     })
  }
  /// Add a function to the module with the name given.
  pub fn add_function<'a>(&'a self, name: &str, sig: &'a Type) -> &'a Function
  {
    let c_name = CString::new(name).unwrap();
    unsafe { core::LLVMAddFunction(self.into(), c_name.as_ptr(), sig.into()) }.into()
  }

  /// Returns the function with the name given, or `None` if no function with that name exists.
  pub fn get_function<'a>(&'a self, name: &str) -> Option<&'a Function>
  {
    let c_name = CString::new(name).unwrap();
    unsafe {
      let ty = core::LLVMGetNamedFunction(self.into(), c_name.as_ptr());
      util::ptr_to_null(ty)
    }
  }

  /// Returns the type with the name given, or `None`` if no type with that name exists.
  pub fn get_type<'a>(&'a self, name: &str) -> Option<&'a Type>
  {
    let c_name = CString::new(name).unwrap();
    unsafe {
      let ty = core::LLVMGetTypeByName(self.into(), c_name.as_ptr());
      util::ptr_to_null(ty)
    }
  }

  /// Clone this module.
  pub fn clone<'a>(&'a self) -> CSemiBox<'a, Module>
  {
    CSemiBox::new(unsafe { core::LLVMCloneModule(self.into()) })
  }

  /// Optimize this module with the given optimization level and size level.
  ///
  /// This runs passes depending on the levels given.
  pub fn optimize(&self, opt_level: usize, size_level: usize) {
    unsafe {
      let builder = builder::LLVMPassManagerBuilderCreate();
      builder::LLVMPassManagerBuilderSetOptLevel(builder, opt_level as c_uint);
      builder::LLVMPassManagerBuilderSetSizeLevel(builder, size_level as c_uint);
      let pass_manager = core::LLVMCreatePassManager();
      builder::LLVMPassManagerBuilderPopulateModulePassManager(builder, pass_manager);
      builder::LLVMPassManagerBuilderDispose(builder);
      core::LLVMRunPassManager(pass_manager, self.into());
    }
  }

  /// Returns the target data of this module represented as a string
  pub fn get_target(&self) -> &str
  {
    unsafe {
      let target = core::LLVMGetTarget(self.into());
      util::to_str(target as *mut c_char)
    }
  }

  pub fn get_data_layout(&self) -> &str
  {
  	unsafe {
  		let layout = core::LLVMGetDataLayout(self.into());
  		util::to_str(layout as *mut c_char)
  	}
  }

  /// Set the target data of this module to the target data string given.
  pub fn set_target(&self, target: &str)
  {
      let c_target = CString::new(target).unwrap();
      unsafe { core::LLVMSetTarget(self.into(), c_target.as_ptr()) }
  }

  /// Verify that the module is safe to run, returning a string detailing the error
  /// when an error occurs.
  pub fn verify(&self) -> Result<(), CBox<str>>
  {
    unsafe {
      let mut error = mem::uninitialized();
      let action = LLVMVerifierFailureAction::LLVMReturnStatusAction;
      if analysis::LLVMVerifyModule(self.into(), action, &mut error) == 1 {
          Err(CBox::new(error))
      } else {
          Ok(())
      }
    }
  }

  /// Compile the module into an object file at the given location.
  ///
  /// Note that this uses the LLVM tool `llc` to do this, which may or may not be
  /// installed on the user's machine.
  pub fn compile(&self, path: &Path, opt_level: usize) -> IoResult<()>
  {
    let dir = env::temp_dir();
    let path = path.to_str().unwrap();
    let mod_path = dir.join("module.bc");
    let mod_path = mod_path.to_str().unwrap();
    try!(self.write_bitcode(mod_path));
    Command::new("llc")
      .arg(&format!("-O={}", opt_level))
      .arg("-filetype=obj")
      .arg("-o").arg(path)
      .arg(mod_path)
      .spawn()
      .map(|_| ())
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does not* destroy the source module.
  pub fn link(&self, src: &Module) -> Result<(), CBox<str>>
  {
    unsafe {
      let dest = self.into();
      let src = src.into();
      let mut message = mem::uninitialized();

      let ret = linker::LLVMLinkModules(dest, src,
      	                                linker::LLVMLinkerMode::LLVMLinkerPreserveSource,
      	                                &mut message);
      if ret == 1 {
        Err(CBox::new(message))
      } else {
        Ok(())
      }
    }
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does* destroy the source module.
  pub fn link_destroy(&self, src: CSemiBox<Module>) -> Result<(), CBox<str>>
  {
    unsafe {
      let dest = self.into();
      let src = src.as_ptr();
      let mut message = mem::uninitialized();

      let ret = linker::LLVMLinkModules(dest, src,
      	                                linker::LLVMLinkerMode::LLVMLinkerDestroySource,
      	                                &mut message);
      if ret == 1 {
        Err(CBox::new(message))
      } else {
        Ok(())
      }
    }
  }

  /// Dump the module to stderr (for debugging).
  pub fn dump(&self)
  {
  	unsafe { core::LLVMDumpModule(self.into()); }
  }
}

get_context!(Module, LLVMGetModuleContext);
impl_display!(Module, LLVMPrintModuleToString);
dispose!(Module, LLVMModule, core::LLVMDisposeModule);

impl<'a> IntoIterator for &'a Module
{
  type Item = &'a Function;
  type IntoIter = Functions<'a>;
  /// Iterate through the functions in the module
  fn into_iter(self) -> Functions<'a>
  {
    Functions {
      value: unsafe { core::LLVMGetFirstFunction(self.into()) },
      marker: PhantomData
    }
  }
}

#[derive(Copy, Clone)]
/// An iterator through the functions contained in a module.
pub struct Functions<'a>
{
  value: LLVMValueRef,
  marker: PhantomData<&'a ()>
}

impl<'a> Iterator for Functions<'a>
{
  type Item = &'a Function;

  fn next(&mut self) -> Option<&'a Function>
  {
    if self.value.is_null() {
      None
    } else {
      let c_next = unsafe { core::LLVMGetNextFunction(self.value) };
      self.value = c_next;
      Some(self.value.into())
    }
  }
}

#[repr(C)]
#[derive(Copy, Clone)]
pub enum AddressSpace
{
  Generic = 0,
  Global = 1,
  Shared = 3,
  Const = 4,
  Local = 5,
}<commit_msg>Fixed a bug which skip the first function in Module.<commit_after>use std::env;
use std::ffi::CString;
use std::fmt;
use std::marker::PhantomData;
use std::mem;
use std::io::{Error, ErrorKind};
use std::io::Result as IoResult;
use std::iter::{Iterator, IntoIterator};
use std::path::Path;
use std::process::Command;

use cbox::{CBox, CSemiBox};
use ffi::analysis::LLVMVerifierFailureAction;
use ffi::{analysis, core, linker, LLVMModule};
use ffi::bit_writer as writer;
use ffi::bit_reader as reader;
use ffi::ir_reader as ir_reader;
use ffi::prelude::{LLVMValueRef, LLVMModuleRef};
use ffi::transforms::pass_manager_builder as builder;
use libc::{c_char, c_uint};

use buffer::MemoryBuffer;
use context::{Context, GetContext};
use util;
use ty::Type;
use value::{Function, GlobalValue, Value};

/// Represents a single compilation unit of code.
///
/// This is attached to the lifetime of the context that constructs it, but is owned by the `CSemiBox`.
pub struct Module;
native_ref!(&Module = LLVMModuleRef);
impl Module
{
  /// Create a new module in the context given with the name given.
  ///
  /// The lifetime of the module will match the lifetime of the context
  /// you instance it in because the context contains it.
  ///
  /// ```rust
  /// use llvm::*;
  /// let context = Context::new();
  /// let context = context.as_semi();
  /// let module = Module::new("name", &context);
  /// println!("{:?}", module)
  /// ```
  pub fn new<'a>(name: &str, context: &'a Context) -> CSemiBox<'a, Module>
  {
    let c_name = CString::new(name).unwrap();
    unsafe {
    	CSemiBox::new(core::LLVMModuleCreateWithNameInContext(c_name.as_ptr(), context.into()))
    }
  }

  /// Add an external global to the module with the given type and name.
  pub fn add_global<'a>(&'a self, name: &str, ty: &'a Type) -> &'a GlobalValue {
    util::with_cstr(name, |ptr| unsafe {
        core::LLVMAddGlobal(self.into(), ty.into(), ptr).into()
    })
  }

  /// Add a global in the given address space to the module with the given type and name.
  pub fn add_global_in_addr_space<'a>(&'a self,
  	                                  name: &str,
  	                                  ty: &'a Type,
  	                                  sp: AddressSpace) -> &'a GlobalValue
  {
    util::with_cstr(name, |ptr| unsafe {
      core::LLVMAddGlobalInAddressSpace(self.into(), ty.into(), ptr, sp as c_uint).into()
    })
  }

  /// Add a constant global to the module with the given type, name and value.
  pub fn add_global_constant<'a>(&'a self, name: &str, val: &'a Value) -> &'a GlobalValue
  {
    util::with_cstr(name, |ptr| unsafe {
      let global = core::LLVMAddGlobal(self.into(), val.get_type().into(), ptr);
      core::LLVMSetInitializer (global.into(), val.into());
      global.into()
    })
  }
  /// Get the global with the name given, or `None` if no global with that name exists.
  pub fn get_global<'a>(&'a self, name: &str) -> Option<&'a GlobalValue>
  {
    util::with_cstr(name, |ptr| unsafe {
        let ptr = core::LLVMGetNamedGlobal(self.into(), ptr);
        util::ptr_to_null(ptr)
    })
  }

  pub fn parse_ir<'a>(context: &'a Context, path: &str) -> Result<CSemiBox<'a, Module>, CBox<str>>
  {
  	unsafe {
      let mut out = mem::uninitialized();
      let mut err = mem::uninitialized();
      let buf = try!(MemoryBuffer::new_from_file(path));

      let ret = ir_reader::LLVMParseIRInContext(context.into(), buf.as_ptr(), &mut out, &mut err);
      if ret == 1 {
      	Err(CBox::new(err))
      } else {
        Ok(CSemiBox::new(out))
      }
    }
  }

  /// Parse this bitcode file into a module, or return an error string.
  pub fn parse_bitcode<'a>(context: &'a Context, path: &str) -> Result<CSemiBox<'a, Module>, CBox<str>> {
    unsafe {
      let mut out = mem::uninitialized();
      let mut err = mem::uninitialized();
      let buf = try!(MemoryBuffer::new_from_file(path));

      let ret = reader::LLVMParseBitcodeInContext(context.into(),
      																					  buf.as_ptr(),
      	                                          &mut out,
      	                                          &mut err);
      if ret == 1 {
          Err(CBox::new(err))
      } else {
          Ok(CSemiBox::new(out))
      }
    }
  }
  /// Write this module's bitcode to the path given.
  pub fn write_bitcode(&self, path: &str) -> IoResult<()>
  {
     util::with_cstr(path, |cpath| unsafe {
       if writer::LLVMWriteBitcodeToFile(self.into(), cpath) != 0 {
           Err(Error::new(ErrorKind::Other, &format!("could not write to {}", path) as &str))
       } else {
           Ok(())
       }
     })
  }
  /// Add a function to the module with the name given.
  pub fn add_function<'a>(&'a self, name: &str, sig: &'a Type) -> &'a Function
  {
    let c_name = CString::new(name).unwrap();
    unsafe { core::LLVMAddFunction(self.into(), c_name.as_ptr(), sig.into()) }.into()
  }

  /// Returns the function with the name given, or `None` if no function with that name exists.
  pub fn get_function<'a>(&'a self, name: &str) -> Option<&'a Function>
  {
    let c_name = CString::new(name).unwrap();
    unsafe {
      let ty = core::LLVMGetNamedFunction(self.into(), c_name.as_ptr());
      util::ptr_to_null(ty)
    }
  }

  /// Returns the type with the name given, or `None`` if no type with that name exists.
  pub fn get_type<'a>(&'a self, name: &str) -> Option<&'a Type>
  {
    let c_name = CString::new(name).unwrap();
    unsafe {
      let ty = core::LLVMGetTypeByName(self.into(), c_name.as_ptr());
      util::ptr_to_null(ty)
    }
  }

  /// Clone this module.
  pub fn clone<'a>(&'a self) -> CSemiBox<'a, Module>
  {
    CSemiBox::new(unsafe { core::LLVMCloneModule(self.into()) })
  }

  /// Optimize this module with the given optimization level and size level.
  ///
  /// This runs passes depending on the levels given.
  pub fn optimize(&self, opt_level: usize, size_level: usize) {
    unsafe {
      let builder = builder::LLVMPassManagerBuilderCreate();
      builder::LLVMPassManagerBuilderSetOptLevel(builder, opt_level as c_uint);
      builder::LLVMPassManagerBuilderSetSizeLevel(builder, size_level as c_uint);
      let pass_manager = core::LLVMCreatePassManager();
      builder::LLVMPassManagerBuilderPopulateModulePassManager(builder, pass_manager);
      builder::LLVMPassManagerBuilderDispose(builder);
      core::LLVMRunPassManager(pass_manager, self.into());
    }
  }

  /// Returns the target data of this module represented as a string
  pub fn get_target(&self) -> &str
  {
    unsafe {
      let target = core::LLVMGetTarget(self.into());
      util::to_str(target as *mut c_char)
    }
  }

  pub fn get_data_layout(&self) -> &str
  {
  	unsafe {
  		let layout = core::LLVMGetDataLayout(self.into());
  		util::to_str(layout as *mut c_char)
  	}
  }

  /// Set the target data of this module to the target data string given.
  pub fn set_target(&self, target: &str)
  {
      let c_target = CString::new(target).unwrap();
      unsafe { core::LLVMSetTarget(self.into(), c_target.as_ptr()) }
  }

  /// Verify that the module is safe to run, returning a string detailing the error
  /// when an error occurs.
  pub fn verify(&self) -> Result<(), CBox<str>>
  {
    unsafe {
      let mut error = mem::uninitialized();
      let action = LLVMVerifierFailureAction::LLVMReturnStatusAction;
      if analysis::LLVMVerifyModule(self.into(), action, &mut error) == 1 {
          Err(CBox::new(error))
      } else {
          Ok(())
      }
    }
  }

  /// Compile the module into an object file at the given location.
  ///
  /// Note that this uses the LLVM tool `llc` to do this, which may or may not be
  /// installed on the user's machine.
  pub fn compile(&self, path: &Path, opt_level: usize) -> IoResult<()>
  {
    let dir = env::temp_dir();
    let path = path.to_str().unwrap();
    let mod_path = dir.join("module.bc");
    let mod_path = mod_path.to_str().unwrap();
    try!(self.write_bitcode(mod_path));
    Command::new("llc")
      .arg(&format!("-O={}", opt_level))
      .arg("-filetype=obj")
      .arg("-o").arg(path)
      .arg(mod_path)
      .spawn()
      .map(|_| ())
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does not* destroy the source module.
  pub fn link(&self, src: &Module) -> Result<(), CBox<str>>
  {
    unsafe {
      let dest = self.into();
      let src = src.into();
      let mut message = mem::uninitialized();

      let ret = linker::LLVMLinkModules(dest, src,
      	                                linker::LLVMLinkerMode::LLVMLinkerPreserveSource,
      	                                &mut message);
      if ret == 1 {
        Err(CBox::new(message))
      } else {
        Ok(())
      }
    }
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does* destroy the source module.
  pub fn link_destroy(&self, src: CSemiBox<Module>) -> Result<(), CBox<str>>
  {
    unsafe {
      let dest = self.into();
      let src = src.as_ptr();
      let mut message = mem::uninitialized();

      let ret = linker::LLVMLinkModules(dest, src,
      	                                linker::LLVMLinkerMode::LLVMLinkerDestroySource,
      	                                &mut message);
      if ret == 1 {
        Err(CBox::new(message))
      } else {
        Ok(())
      }
    }
  }

  /// Dump the module to stderr (for debugging).
  pub fn dump(&self)
  {
  	unsafe { core::LLVMDumpModule(self.into()); }
  }
}

get_context!(Module, LLVMGetModuleContext);
impl_display!(Module, LLVMPrintModuleToString);
dispose!(Module, LLVMModule, core::LLVMDisposeModule);

impl<'a> IntoIterator for &'a Module
{
  type Item = &'a Function;
  type IntoIter = Functions<'a>;
  /// Iterate through the functions in the module
  fn into_iter(self) -> Functions<'a>
  {
    Functions {
      value: unsafe { core::LLVMGetFirstFunction(self.into()) },
      marker: PhantomData
    }
  }
}

#[derive(Copy, Clone)]
/// An iterator through the functions contained in a module.
pub struct Functions<'a>
{
  value: LLVMValueRef,
  marker: PhantomData<&'a ()>
}

impl<'a> Iterator for Functions<'a>
{
  type Item = &'a Function;

  fn next(&mut self) -> Option<&'a Function>
  {
    let old: LLVMValueRef = self.value;
    if !old.is_null() {
       self.value = unsafe { core::LLVMGetNextFunction(old) };
       Some(old.into())
    } else {
       None
    }
  }
}

#[repr(C)]
#[derive(Copy, Clone)]
pub enum AddressSpace
{
  Generic = 0,
  Global = 1,
  Shared = 3,
  Const = 4,
  Local = 5,
}<|endoftext|>


<commit_before>//#![allow(unused)]

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex, Arc};
use std::vec::Vec;
use std::collections::HashMap;
use std::cell::UnsafeCell;
use std::hash::{Hash, Hasher};
use std::mem::{align_of, size_of_val, transmute};
use std::slice::bytes::copy_memory;

use ::Result;

pub struct GlobalState { version : AtomicUsize }

/// Baseline implementation of GlobalState, currently do not hqve the aid of the compiler, so this
/// may have major issues.
/// TODO: make a proper implementation of the GlobalState trait
impl GlobalState {
    fn newLocal<'l, 'g>(&'g self) -> LocalState<'l, 'g> {
        LocalState {
            reads: Vec::new(),
            writes: HashMap::new(),
            global: self,
            snapshot: self.version.load(Ordering::SeqCst),
        }
    }

    pub fn retry<A, F>(&self, l: &mut LocalState, mut func: F) -> A
    where F: FnMut(&mut LocalState) -> Result<A> {
        let bx = Arc::new((Mutex::new(false), Condvar::new()));
        for &(addr, size) in &l.reads {
            let mut waiters = addr.addr.waiters.lock().unwrap();
            // this is actually unsafe...
            // maybe protect waiters with a lock
            waiters.push(bx.clone());
            //addr.waiters().push((&mtx, &cv));
        }
        let (ref mtx, ref cv) = *bx;
        let mut done = mtx.lock().unwrap();
        while !*done {
            done = cv.wait(done).unwrap();
        }
        self.run(func)
    }

    pub fn run<A, F>(&self, mut func : F) -> A
    where F: FnMut(&mut LocalState) -> Result<A> {
        let mut local = self.newLocal();
        match func(&mut local).spec_move() {
            Ok(r) => r,
            _ => self.run(func)
        }
    }

    pub fn or_else<A, F, G>(&self, mut func1 : F, func2 : G) -> A
        where F: FnMut(&mut LocalState) -> Result<A>,
              G: FnMut(&mut LocalState) -> Result<A> {
            let mut local1 = self.newLocal();
            match func1(&mut local1).spec_move() {
                Ok(r) => r,
                //TODO: this may not work, compiler may not know that
                //local1, which borrows self, has relinquised ownership
                _ => self.run(func2)
            }
    }
}

pub struct LocalState<'v, 'g : 'v> {
    /// the norec paper has the reads stored as a list of <address, value> pairs. For us it
    /// would probably be better to have versioned TVars. That way we always have a fast
    /// comparison operation and avoid doing extra dynamic dispatch when validating reads.
    reads : Vec<(TVarAddr<'v, 'g>, usize)>,

    /// for writes we will probably have to cast unsafely to get this to work. TVars just hold
    /// arbitrary types and we only have one table.  XXX: gross
    writes : HashMap<TVarAddr<'v, 'g>, Box<[u8]>>,

    /// STM environment we are using
    /// TODO: how are we getting away with this not being a mutable reference?
    global : &'g GlobalState,

    /// last version at which our view of the global state was valid
    snapshot : usize
}

impl<'v, 'g : 'v> LocalState<'v, 'g> {

    /// Validate the current local state against the state of the world,
    /// Aborts if a TVar in our readset has been written, otherwise
    /// returns the version number to with we are consistent
    pub fn validate(&self) -> Result<usize> {
        loop {
            let time = self.global.version.load(Ordering::SeqCst);
            if (time & 1) == 0 {
                continue;
            }
            for &(addr, version) in &self.reads {
                if addr.version().load(Ordering::SeqCst) != version {
                    return Result::abort();
                }
            }
            return Result::ret(time);
        }
    }

    /// Acquire reference to value in TVar within the transaction
    pub fn read<'l, A>(&'l mut self, tvar : &'v TVar<'g, A>) -> Result<&'l A>
        where 'v : 'l
    {
        let addr: TVarAddr<'v, 'g> = unsafe { tvar.to_addr() } ;

        // Should be &&'l ?
        if let Some(v) = self.writes.get(&addr) {
            let r : && [u8] = &&**v;
            assert_eq!(r.len(), ::std::mem::size_of::<A>());
            let r2 : && A = unsafe { transmute(r) };
            return Result::ret(r2)
        }

        if self.snapshot == self.global.version.load(Ordering::SeqCst) {
            self.reads.push((addr, tvar.version.load(Ordering::SeqCst)));
            Result::ret(unsafe {tvar.get_val()})
        } else {
            self.validate().bind( | v | {
                assert!(self.snapshot <= v);
                self.snapshot = v;
                self.reads.push((addr, tvar.version.load(Ordering::SeqCst)));
                Result::ret(unsafe {tvar.get_val()})
            })
        }
    }

    /// Add the TVar and value to our writeset, these are thread-local until the
    /// transaction commits
    pub fn write<A : Copy>(&mut self, tvar : &'v mut TVar<'g, A>, val : A) -> Result<()> {
        let b : Box<[A]> = Box::new([val]);
        unsafe {
            self.writes.insert(tvar.to_addr(), transmute(b));
        }
        Result::ret(())
    }

    /// if (read−only transaction)
    ///     return
    /// while (!CAS(&global lock, snapshot, snapshot + 1))
    ///     snapshot = Validate()
    /// for each (addr, val) in writes
    ///     ∗addr = val
    /// global lock = snapshot + 2 // one more than CAS above
    pub fn commit(&mut self) -> Result<()> {
        fn commit_loop(l : &mut LocalState) -> Result<()> {
            if l.global
                .version
                .compare_and_swap(l.snapshot, l.snapshot + 1, Ordering::SeqCst) == l.snapshot {
                Result::ret(())
            } else {
                l.validate().bind(|v| {
                    l.snapshot = v;
                    commit_loop(l)
                })
            }
        }
        if self.writes.is_empty() {
            Result::ret(())
        } else {
            commit_loop(self).bind(| _ | {
                for (mut addr, value) in &self.writes {
                    let dst = unsafe { addr.aligned_mem() };
                    debug_assert_eq!(value.len(), dst.len());
                    copy_memory(value, dst);
                    addr.addr.version.fetch_add(1, Ordering::SeqCst);
                    //TODO: DCLP if this is too slow in the common case that waiters is empty
                    let mut waiters = addr.addr.waiters.lock().unwrap();
                    for bx in &*waiters {
                        let (ref mtx, ref cv) = **bx;
                        let mut done = mtx.lock().unwrap();
                        *done = true;
                        cv.notify_all();
                    }
                    // this is safe because we hold the sequence lock
                    waiters.clear();
                    //addr.waiters().clear();
                    //TODO: set waiters to empty
                }
                self.global
                    .version
                    .compare_and_swap(self.snapshot + 1, self.snapshot + 2, Ordering::SeqCst);
                Result::ret(())
            })
        }
    }
}

/// Haskell-style container of mutable state within a transaction
/// TODO: impl TVar trait
pub struct TVar<'a, A: ?Sized> {
    //TODO add in mutex/cv pair for waiting for explicit aborts
    //This includes an explicit abort() command (as opposed to failing validation,
    //which should cause an immediate retry).
    version : AtomicUsize,
    //TODO this can be of lifetime 'b : 'a
    waiters : Mutex<Vec<Arc<(Mutex<bool>, Condvar)>>>,
    global : &'a GlobalState,
    val : UnsafeCell<A>,
}

impl<'a, A: ?Sized> TVar<'a, A> {
    unsafe fn get_val(&self) -> &A {
        &*self.val.get()
    }

    unsafe fn get_val_mut(&self) -> &mut A {
        &mut *self.val.get()
    }
}

impl<'a, A> TVar<'a, A> {
    unsafe fn to_addr<'b>(&'b self) -> TVarAddr<'b, 'a> {
        use std::raw::Slice;
        let fat_ptr = Slice::<u8> {
            data: transmute(self),
            len: size_of_val(self),
        };
        //let ptr : &*mut TVar<'a, [u8]> = transmute(self));
        TVarAddr {
            align: align_of::<A>(),
            addr: transmute(fat_ptr),
        }
    }

    ///TODO: change trait API to use this signature
    fn read<'l, 'b> (&'b mut self, l: &'l mut LocalState<'b, 'a>) -> Result<&'l A> {
        l.read(self)
    }

}
impl<'a, A : Copy> TVar<'a, A> {
    ///TODO: change trait API to use this signature
    fn write<'l, 'b>(&'b mut self, l: &'l mut LocalState<'b, 'a>, arg : A) -> Result<()> {
        l.write(self, arg)
    }
}

#[derive(Copy, Clone)]
/// Lies... Lies
struct TVarAddr<'a, 'b : 'a> {
    align: usize,
    addr: &'a TVar<'b, [u8]>, //*mut
    //_phan: PhantomData<&'a TVar<'b, ()>>,
}

impl<'a, 'b : 'a> Hash for TVarAddr<'a, 'b> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.raw_addr().hash(state);
    }
}

impl<'a, 'b : 'a> Eq for TVarAddr<'a, 'b> {}

impl<'a, 'b : 'a> PartialEq for TVarAddr<'a, 'b> {
    fn eq(&self, other: &Self) -> bool {
        let c = self.raw_addr() == other.raw_addr();
        if c {
            debug_assert_eq!(self.raw_len(), other.raw_len());
            debug_assert_eq!(self.align, other.align);
        }
        c
    }
}

impl<'a, 'b : 'a> TVarAddr<'a, 'b> {
    fn raw_addr(self) -> usize {
        self.addr as *const _ as *const () as usize
    }

    fn raw_len(self) -> usize {
        let y = self.addr.val.get();
        unsafe { transmute::<_,&[u8]>(y).len() }
    }

    unsafe fn aligned_mem(self) -> &'a mut [u8] {
        let buf = self.addr.val.get();
        let p = buf as *mut () as usize;
        // Assumes alignment is power of 2
        let o = ((p + (self.align - 1)) & self.align) - p;
        &mut (*buf)[o..]
    }

    fn version(self) -> &'a AtomicUsize {
        &(*self.addr).version
    }

    fn waiters(self) -> &'a mut Vec<(&'b Mutex<bool>, &'b Condvar)> {
        panic!("nyi")
    }
}

pub fn align_down_mut<T>(sp: *mut T, n: usize) -> *mut T {
  let sp = (sp as usize) & !(n - 1);
  sp as *mut T
}
<commit_msg>removed irrelevant comment<commit_after>//#![allow(unused)]

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex, Arc};
use std::vec::Vec;
use std::collections::HashMap;
use std::cell::UnsafeCell;
use std::hash::{Hash, Hasher};
use std::mem::{align_of, size_of_val, transmute};
use std::slice::bytes::copy_memory;

use ::Result;

pub struct GlobalState { version : AtomicUsize }

/// Baseline implementation of GlobalState, currently do not hqve the aid of the compiler, so this
/// may have major issues.
/// TODO: make a proper implementation of the GlobalState trait
impl GlobalState {
    fn newLocal<'l, 'g>(&'g self) -> LocalState<'l, 'g> {
        LocalState {
            reads: Vec::new(),
            writes: HashMap::new(),
            global: self,
            snapshot: self.version.load(Ordering::SeqCst),
        }
    }

    pub fn retry<A, F>(&self, l: &mut LocalState, mut func: F) -> A
    where F: FnMut(&mut LocalState) -> Result<A> {
        let bx = Arc::new((Mutex::new(false), Condvar::new()));
        for &(addr, size) in &l.reads {
            let mut waiters = addr.addr.waiters.lock().unwrap();
            // this is actually unsafe...
            // maybe protect waiters with a lock
            waiters.push(bx.clone());
            //addr.waiters().push((&mtx, &cv));
        }
        let (ref mtx, ref cv) = *bx;
        let mut done = mtx.lock().unwrap();
        while !*done {
            done = cv.wait(done).unwrap();
        }
        self.run(func)
    }

    pub fn run<A, F>(&self, mut func : F) -> A
        where F: FnMut(&mut LocalState) -> Result<A> {
            let mut local = self.newLocal();
            match func(&mut local).spec_move() {
                Ok(r) => r,
                _ => self.run(func)
            }
    }

    pub fn or_else<A, F, G>(&self, mut func1 : F, func2 : G) -> A
        where F: FnMut(&mut LocalState) -> Result<A>,
              G: FnMut(&mut LocalState) -> Result<A> {
            let mut local1 = self.newLocal();
            match func1(&mut local1).spec_move() {
                Ok(r) => r,
                //TODO: this may not work, compiler may not know that
                //local1, which borrows self, has relinquised ownership
                _ => self.run(func2)
            }
    }
}

pub struct LocalState<'v, 'g : 'v> {
    /// the norec paper has the reads stored as a list of <address, value> pairs. For us it
    /// would probably be better to have versioned TVars. That way we always have a fast
    /// comparison operation and avoid doing extra dynamic dispatch when validating reads.
    reads : Vec<(TVarAddr<'v, 'g>, usize)>,

    /// for writes we will probably have to cast unsafely to get this to work. TVars just hold
    /// arbitrary types and we only have one table.  XXX: gross
    writes : HashMap<TVarAddr<'v, 'g>, Box<[u8]>>,

    /// STM environment we are using
    /// TODO: how are we getting away with this not being a mutable reference?
    global : &'g GlobalState,

    /// last version at which our view of the global state was valid
    snapshot : usize
}

impl<'v, 'g : 'v> LocalState<'v, 'g> {

    /// Validate the current local state against the state of the world,
    /// Aborts if a TVar in our readset has been written, otherwise
    /// returns the version number to with we are consistent
    pub fn validate(&self) -> Result<usize> {
        loop {
            let time = self.global.version.load(Ordering::SeqCst);
            if (time & 1) == 0 {
                continue;
            }
            for &(addr, version) in &self.reads {
                if addr.version().load(Ordering::SeqCst) != version {
                    return Result::abort();
                }
            }
            return Result::ret(time);
        }
    }

    /// Acquire reference to value in TVar within the transaction
    pub fn read<'l, A>(&'l mut self, tvar : &'v TVar<'g, A>) -> Result<&'l A>
        where 'v : 'l
    {
        let addr: TVarAddr<'v, 'g> = unsafe { tvar.to_addr() } ;

        // Should be &&'l ?
        if let Some(v) = self.writes.get(&addr) {
            let r : && [u8] = &&**v;
            assert_eq!(r.len(), ::std::mem::size_of::<A>());
            let r2 : && A = unsafe { transmute(r) };
            return Result::ret(r2)
        }

        if self.snapshot == self.global.version.load(Ordering::SeqCst) {
            self.reads.push((addr, tvar.version.load(Ordering::SeqCst)));
            Result::ret(unsafe {tvar.get_val()})
        } else {
            self.validate().bind( | v | {
                assert!(self.snapshot <= v);
                self.snapshot = v;
                self.reads.push((addr, tvar.version.load(Ordering::SeqCst)));
                Result::ret(unsafe {tvar.get_val()})
            })
        }
    }

    /// Add the TVar and value to our writeset, these are thread-local until the
    /// transaction commits
    pub fn write<A : Copy>(&mut self, tvar : &'v mut TVar<'g, A>, val : A) -> Result<()> {
        let b : Box<[A]> = Box::new([val]);
        unsafe {
            self.writes.insert(tvar.to_addr(), transmute(b));
        }
        Result::ret(())
    }

    /// if (read−only transaction)
    ///     return
    /// while (!CAS(&global lock, snapshot, snapshot + 1))
    ///     snapshot = Validate()
    /// for each (addr, val) in writes
    ///     ∗addr = val
    /// global lock = snapshot + 2 // one more than CAS above
    pub fn commit(&mut self) -> Result<()> {
        fn commit_loop(l : &mut LocalState) -> Result<()> {
            if l.global
                .version
                .compare_and_swap(l.snapshot, l.snapshot + 1, Ordering::SeqCst) == l.snapshot {
                Result::ret(())
            } else {
                l.validate().bind(|v| {
                    l.snapshot = v;
                    commit_loop(l)
                })
            }
        }
        if self.writes.is_empty() {
            Result::ret(())
        } else {
            commit_loop(self).bind(| _ | {
                for (mut addr, value) in &self.writes {
                    let dst = unsafe { addr.aligned_mem() };
                    debug_assert_eq!(value.len(), dst.len());
                    copy_memory(value, dst);
                    addr.addr.version.fetch_add(1, Ordering::SeqCst);
                    //TODO: DCLP if this is too slow in the common case that waiters is empty
                    let mut waiters = addr.addr.waiters.lock().unwrap();
                    for bx in &*waiters {
                        let (ref mtx, ref cv) = **bx;
                        let mut done = mtx.lock().unwrap();
                        *done = true;
                        cv.notify_all();
                    }
                    waiters.clear();
                }
                self.global
                    .version
                    .compare_and_swap(self.snapshot + 1, self.snapshot + 2, Ordering::SeqCst);
                Result::ret(())
            })
        }
    }
}

/// Haskell-style container of mutable state within a transaction
/// TODO: impl TVar trait
pub struct TVar<'a, A: ?Sized> {
    //TODO add in mutex/cv pair for waiting for explicit aborts
    //This includes an explicit abort() command (as opposed to failing validation,
    //which should cause an immediate retry).
    version : AtomicUsize,
    //TODO this can be of lifetime 'b : 'a
    waiters : Mutex<Vec<Arc<(Mutex<bool>, Condvar)>>>,
    global : &'a GlobalState,
    val : UnsafeCell<A>,
}

impl<'a, A: ?Sized> TVar<'a, A> {
    unsafe fn get_val(&self) -> &A {
        &*self.val.get()
    }

    unsafe fn get_val_mut(&self) -> &mut A {
        &mut *self.val.get()
    }
}

impl<'a, A> TVar<'a, A> {
    unsafe fn to_addr<'b>(&'b self) -> TVarAddr<'b, 'a> {
        use std::raw::Slice;
        let fat_ptr = Slice::<u8> {
            data: transmute(self),
            len: size_of_val(self),
        };
        //let ptr : &*mut TVar<'a, [u8]> = transmute(self));
        TVarAddr {
            align: align_of::<A>(),
            addr: transmute(fat_ptr),
        }
    }

    ///TODO: change trait API to use this signature
    fn read<'l, 'b> (&'b mut self, l: &'l mut LocalState<'b, 'a>) -> Result<&'l A> {
        l.read(self)
    }

}
impl<'a, A : Copy> TVar<'a, A> {
    ///TODO: change trait API to use this signature
    fn write<'l, 'b>(&'b mut self, l: &'l mut LocalState<'b, 'a>, arg : A) -> Result<()> {
        l.write(self, arg)
    }
}

#[derive(Copy, Clone)]
/// Lies... Lies
struct TVarAddr<'a, 'b : 'a> {
    align: usize,
    addr: &'a TVar<'b, [u8]>, //*mut
    //_phan: PhantomData<&'a TVar<'b, ()>>,
}

impl<'a, 'b : 'a> Hash for TVarAddr<'a, 'b> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.raw_addr().hash(state);
    }
}

impl<'a, 'b : 'a> Eq for TVarAddr<'a, 'b> {}

impl<'a, 'b : 'a> PartialEq for TVarAddr<'a, 'b> {
    fn eq(&self, other: &Self) -> bool {
        let c = self.raw_addr() == other.raw_addr();
        if c {
            debug_assert_eq!(self.raw_len(), other.raw_len());
            debug_assert_eq!(self.align, other.align);
        }
        c
    }
}

impl<'a, 'b : 'a> TVarAddr<'a, 'b> {
    fn raw_addr(self) -> usize {
        self.addr as *const _ as *const () as usize
    }

    fn raw_len(self) -> usize {
        let y = self.addr.val.get();
        unsafe { transmute::<_,&[u8]>(y).len() }
    }

    unsafe fn aligned_mem(self) -> &'a mut [u8] {
        let buf = self.addr.val.get();
        let p = buf as *mut () as usize;
        // Assumes alignment is power of 2
        let o = ((p + (self.align - 1)) & self.align) - p;
        &mut (*buf)[o..]
    }

    fn version(self) -> &'a AtomicUsize {
        &(*self.addr).version
    }

    fn waiters(self) -> &'a mut Vec<(&'b Mutex<bool>, &'b Condvar)> {
        panic!("nyi")
    }
}

pub fn align_down_mut<T>(sp: *mut T, n: usize) -> *mut T {
  let sp = (sp as usize) & !(n - 1);
  sp as *mut T
}
<|endoftext|>