<commit_before>use super::{SeasonState, SeriesConfig};
use crate::backend::{AnimeEntry, AnimeInfo, SyncBackend};
use crate::error::SeriesError;
use crate::input::{self, Answer};
use lazy_static::lazy_static;
use regex::Regex;
use serde_derive::{Deserialize, Serialize};
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use toml;

pub struct FolderData {
    pub series_info: SeriesInfo,
    pub savefile: SaveData,
    pub subseries: String,
    pub path: PathBuf,
}

impl FolderData {
    const DEFAULT_SUBSERIES_NAME: &'static str = "main";

    pub fn load_dir(path: &Path, subseries: Option<String>) -> Result<FolderData, SeriesError> {
        let subseries = subseries.unwrap_or_else(|| FolderData::DEFAULT_SUBSERIES_NAME.into());

        let mut savefile = SaveData::from_dir(path)?;
        let series_info = FolderData::load_series_info(path, &mut savefile, &subseries)?;

        Ok(FolderData {
            series_info,
            savefile,
            subseries,
            path: PathBuf::from(path),
        })
    }

    fn load_series_info(
        path: &Path,
        savefile: &mut SaveData,
        subseries: &str,
    ) -> Result<SeriesInfo, SeriesError> {
        let mut ep_data = parse_episode_files_until_valid(path, &mut savefile.episode_matcher)?;
        let subseries_data = savefile.get_mut_subseries_entry(subseries);

        if let Some(info) = SeriesInfo::select_from_subseries(&mut ep_data, &subseries_data)? {
            return Ok(info);
        }

        let info = SeriesInfo::prompt_select_from_episodes(ep_data)?;
        subseries_data.files_title = Some(info.name.clone());

        Ok(info)
    }

    pub fn populate_season_data<B>(&mut self, config: &SeriesConfig<B>) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        let num_seasons = self.seasons().len();

        if num_seasons > config.season_num {
            return Ok(());
        }

        for cur_season in num_seasons..=config.season_num {
            let info = AnimeInfo::default();
            let entry = AnimeEntry::new(info);

            let mut season = SeasonState {
                state: entry,
                needs_info: true,
                needs_sync: config.offline_mode,
            };

            season.sync_info_from_remote(config, &self, cur_season)?;

            let subseries = self
                .savefile
                .get_mut_subseries_entry(self.subseries.clone());

            subseries.season_states.push(season);
        }

        Ok(())
    }

    pub fn calculate_season_offset(&self, range: Range<usize>) -> Option<u32> {
        let seasons = self.seasons();
        let num_seasons = seasons.len();

        if range.end >= num_seasons {
            return None;
        }

        let mut offset = 0;

        for i in range {
            let season = &seasons[i];

            match season.state.info.episodes {
                Some(eps) => offset += eps,
                None => return Some(offset),
            }
        }

        Some(offset)
    }

    pub fn delete_series_dir(&self) -> Result<(), SeriesError> {
        println!("WARNING: {} will be deleted", self.path.to_string_lossy());
        println!("is this ok? (y/N)");

        if !input::read_yn(Answer::No)? {
            return Ok(());
        }

        fs::remove_dir_all(&self.path)?;
        Ok(())
    }

    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.series_info.get_episode(episode)
    }

    pub fn save(&self) -> Result<(), SeriesError> {
        self.savefile.write_to_file()
    }

    pub fn seasons(&self) -> &Vec<SeasonState> {
        &self.savefile.subseries[&self.subseries].season_states
    }

    pub fn seasons_mut(&mut self) -> &mut Vec<SeasonState> {
        &mut self
            .savefile
            .get_mut_subseries_entry(self.subseries.clone())
            .season_states
    }
}

pub type SubSeriesName = String;

#[derive(Serialize, Deserialize)]
pub struct SaveData {
    pub episode_matcher: Option<String>,
    #[serde(flatten)]
    pub subseries: HashMap<SubSeriesName, SubSeriesData>,
    #[serde(skip)]
    pub path: PathBuf,
}

impl SaveData {
    const DATA_FILE_NAME: &'static str = ".anup";

    pub fn new(path: PathBuf) -> SaveData {
        SaveData {
            episode_matcher: None,
            subseries: HashMap::new(),
            path,
        }
    }

    pub fn get_mut_subseries_entry<S>(&mut self, name: S) -> &mut SubSeriesData
    where
        S: Into<String>,
    {
        self.subseries
            .entry(name.into())
            .or_insert_with(SubSeriesData::new)
    }

    pub fn from_dir(path: &Path) -> Result<SaveData, SeriesError> {
        let path = PathBuf::from(path).join(SaveData::DATA_FILE_NAME);

        if !path.exists() {
            return Ok(SaveData::new(path));
        }

        let file_contents = fs::read_to_string(&path)?;

        let mut save_data: SaveData = toml::from_str(&file_contents)?;
        save_data.path = path;

        Ok(save_data)
    }

    pub fn write_to_file(&self) -> Result<(), SeriesError> {
        let toml = toml::to_string(self)?;
        fs::write(&self.path, toml)?;

        Ok(())
    }
}

#[derive(Serialize, Deserialize)]
pub struct SubSeriesData {
    pub files_title: Option<String>,
    pub season_states: Vec<SeasonState>,
}

impl SubSeriesData {
    pub fn new() -> SubSeriesData {
        SubSeriesData {
            files_title: None,
            season_states: Vec::new(),
        }
    }
}

pub type EpisodePaths = HashMap<u32, PathBuf>;
pub type SeriesEpisodes = HashMap<String, EpisodePaths>;

pub struct SeriesInfo {
    pub name: String,
    pub episodes: HashMap<u32, PathBuf>,
}

impl SeriesInfo {
    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.episodes
            .get(&episode)
            .ok_or_else(|| SeriesError::EpisodeNotFound(episode))
    }

    pub fn select_from_subseries(
        ep_data: &mut SeriesEpisodes,
        subseries: &SubSeriesData,
    ) -> Result<Option<SeriesInfo>, SeriesError> {
        if let Some(name) = &subseries.files_title {
            let entry = ep_data.remove_entry(name);

            match entry {
                Some((name, episodes)) => return Ok(Some(SeriesInfo { name, episodes })),
                None => return Err(SeriesError::NoSeriesEpisodes),
            }
        }

        Ok(None)
    }

    pub fn prompt_select_from_episodes(info: SeriesEpisodes) -> Result<SeriesInfo, SeriesError> {
        if info.is_empty() {
            return Err(SeriesError::NoSeriesFound);
        }

        let mut info = info
            .into_iter()
            .map(|(name, eps)| SeriesInfo {
                name,
                episodes: eps,
            })
            .collect::<Vec<_>>();

        if info.len() == 1 {
            return Ok(info.swap_remove(0));
        }

        println!("multiple series found in directory");
        println!("please enter the number next to the episode files you want to use:");

        for (i, series) in info.iter().enumerate() {
            println!("{} [{}]", 1 + i, series.name);
        }

        let index = input::read_range(1, info.len())? - 1;
        let series = info.swap_remove(index);

        Ok(series)
    }
}

struct EpisodeFile {
    series_name: String,
    episode_num: u32,
}

impl EpisodeFile {
    fn parse(path: &Path, matcher: &Regex) -> Result<EpisodeFile, SeriesError> {
        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?;

        let caps = matcher
            .captures(&filename)
            .ok_or(SeriesError::EpisodeRegexCaptureFailed)?;

        let series_name = caps
            .name("name")
            .ok_or_else(|| SeriesError::UnknownRegexCapture("name".into()))
            .map(|name| EpisodeFile::cleanup_title(name.as_str()))?;

        let episode = match caps.name("episode") {
            Some(capture) => capture
                .as_str()
                .parse()
                .map_err(SeriesError::EpisodeNumParseFailed)?,
            None => 1,
        };

        Ok(EpisodeFile {
            series_name,
            episode_num: episode,
        })
    }

    fn cleanup_title(title: &str) -> String {
        let mut clean_title = String::from(title);
        let bytes = unsafe { clean_title.as_bytes_mut() };

        for byte in bytes {
            match byte {
                b'.' | b'_' => *byte = b' ',
                _ => (),
            }
        }

        clean_title
    }
}

// This default pattern will match episodes in several common formats, such as:
// [Group] Series Name - 01.mkv
// [Group]_Series_Name_-_01.mkv
// [Group].Series.Name.-.01.mkv
// [Group] Series Name - 01 [tag 1][tag 2].mkv
// [Group]_Series_Name_-_01_[tag1][tag2].mkv
// [Group].Series.Name.-.01.[tag1][tag2].mkv
// Series Name - 01.mkv
// Series_Name_-_01.mkv
// Series.Name.-.01.mkv
const EP_FORMAT_REGEX: &str =
    r"(?:\[.+?\](?:_+|\.+|\s*))?(?P<name>.+?)(?:\s*|_*|\.*)-(?:\s*|_*|\.*)(?P<episode>\d+)";

fn format_episode_parser_regex<'a, S>(pattern: Option<S>) -> Result<Cow<'a, Regex>, SeriesError>
where
    S: AsRef<str>,
{
    lazy_static! {
        static ref EP_FORMAT: Regex = Regex::new(EP_FORMAT_REGEX).unwrap();
    }

    match pattern {
        Some(pattern) => {
            let pattern = pattern
                .as_ref()
                .replace("{name}", "(?P<name>.+?)")
                .replace("{episode}", r"(?P<episode>\d+)");

            let regex = Regex::new(&pattern)?;
            Ok(Cow::Owned(regex))
        }
        None => Ok(Cow::Borrowed(&*EP_FORMAT)),
    }
}

pub fn parse_episode_files<S>(
    path: &Path,
    pattern: Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str>,
{
    if !path.is_dir() {
        return Err(SeriesError::NotADirectory(path.to_string_lossy().into()));
    }

    let pattern = format_episode_parser_regex(pattern)?;
    let mut data = HashMap::with_capacity(1);

    for entry in fs::read_dir(path).map_err(SeriesError::Io)? {
        let entry = entry.map_err(SeriesError::Io)?.path();

        if !entry.is_file() {
            continue;
        }

        let episode = match EpisodeFile::parse(&entry, pattern.as_ref()) {
            Ok(episode) => episode,
            Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
            Err(err) => return Err(err),
        };

        let series = data
            .entry(episode.series_name)
            .or_insert_with(|| HashMap::with_capacity(1));

        series.insert(episode.episode_num, entry);
    }

    if data.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    Ok(data)
}

pub fn parse_episode_files_until_valid<S>(
    path: &Path,
    pattern: &mut Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str> + From<String>,
{
    loop {
        match parse_episode_files(path, pattern.as_ref()) {
            Ok(data) => return Ok(data),
            Err(SeriesError::NoSeriesFound) => {
                eprintln!("error: no series found");
                println!("you will now be prompted to enter a custom regex pattern");
                println!("when entering the pattern, please mark the series name and episode number with {{name}} and {{episode}}, respectively");
                println!("note: if you're trying to play a single file (like a movie) and it doesn't have an episode number, you can omit the {{episode}} marker");
                println!("example:");
                println!("  filename: [SubGroup] Series Name - Ep01 [1080p].mkv");
                println!(r"  pattern: \[.+?\] {{name}} - Ep{{episode}} \[1080p\].mkv");
                println!("please enter your custom pattern:");
            }
            Err(err @ SeriesError::Regex(_)) | Err(err @ SeriesError::UnknownRegexCapture(_)) => {
                eprintln!("error parsing regex pattern: {}", err);
                println!("please try again:");
            }
            Err(err) => return Err(err),
        }

        let line = input::read_line()?;
        *pattern = Some(line.into());
    }
}
<commit_msg>[series] Improve episode parsing regex<commit_after>use super::{SeasonState, SeriesConfig};
use crate::backend::{AnimeEntry, AnimeInfo, SyncBackend};
use crate::error::SeriesError;
use crate::input::{self, Answer};
use lazy_static::lazy_static;
use regex::Regex;
use serde_derive::{Deserialize, Serialize};
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use toml;

pub struct FolderData {
    pub series_info: SeriesInfo,
    pub savefile: SaveData,
    pub subseries: String,
    pub path: PathBuf,
}

impl FolderData {
    const DEFAULT_SUBSERIES_NAME: &'static str = "main";

    pub fn load_dir(path: &Path, subseries: Option<String>) -> Result<FolderData, SeriesError> {
        let subseries = subseries.unwrap_or_else(|| FolderData::DEFAULT_SUBSERIES_NAME.into());

        let mut savefile = SaveData::from_dir(path)?;
        let series_info = FolderData::load_series_info(path, &mut savefile, &subseries)?;

        Ok(FolderData {
            series_info,
            savefile,
            subseries,
            path: PathBuf::from(path),
        })
    }

    fn load_series_info(
        path: &Path,
        savefile: &mut SaveData,
        subseries: &str,
    ) -> Result<SeriesInfo, SeriesError> {
        let mut ep_data = parse_episode_files_until_valid(path, &mut savefile.episode_matcher)?;
        let subseries_data = savefile.get_mut_subseries_entry(subseries);

        if let Some(info) = SeriesInfo::select_from_subseries(&mut ep_data, &subseries_data)? {
            return Ok(info);
        }

        let info = SeriesInfo::prompt_select_from_episodes(ep_data)?;
        subseries_data.files_title = Some(info.name.clone());

        Ok(info)
    }

    pub fn populate_season_data<B>(&mut self, config: &SeriesConfig<B>) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        let num_seasons = self.seasons().len();

        if num_seasons > config.season_num {
            return Ok(());
        }

        for cur_season in num_seasons..=config.season_num {
            let info = AnimeInfo::default();
            let entry = AnimeEntry::new(info);

            let mut season = SeasonState {
                state: entry,
                needs_info: true,
                needs_sync: config.offline_mode,
            };

            season.sync_info_from_remote(config, &self, cur_season)?;

            let subseries = self
                .savefile
                .get_mut_subseries_entry(self.subseries.clone());

            subseries.season_states.push(season);
        }

        Ok(())
    }

    pub fn calculate_season_offset(&self, range: Range<usize>) -> Option<u32> {
        let seasons = self.seasons();
        let num_seasons = seasons.len();

        if range.end >= num_seasons {
            return None;
        }

        let mut offset = 0;

        for i in range {
            let season = &seasons[i];

            match season.state.info.episodes {
                Some(eps) => offset += eps,
                None => return Some(offset),
            }
        }

        Some(offset)
    }

    pub fn delete_series_dir(&self) -> Result<(), SeriesError> {
        println!("WARNING: {} will be deleted", self.path.to_string_lossy());
        println!("is this ok? (y/N)");

        if !input::read_yn(Answer::No)? {
            return Ok(());
        }

        fs::remove_dir_all(&self.path)?;
        Ok(())
    }

    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.series_info.get_episode(episode)
    }

    pub fn save(&self) -> Result<(), SeriesError> {
        self.savefile.write_to_file()
    }

    pub fn seasons(&self) -> &Vec<SeasonState> {
        &self.savefile.subseries[&self.subseries].season_states
    }

    pub fn seasons_mut(&mut self) -> &mut Vec<SeasonState> {
        &mut self
            .savefile
            .get_mut_subseries_entry(self.subseries.clone())
            .season_states
    }
}

pub type SubSeriesName = String;

#[derive(Serialize, Deserialize)]
pub struct SaveData {
    pub episode_matcher: Option<String>,
    #[serde(flatten)]
    pub subseries: HashMap<SubSeriesName, SubSeriesData>,
    #[serde(skip)]
    pub path: PathBuf,
}

impl SaveData {
    const DATA_FILE_NAME: &'static str = ".anup";

    pub fn new(path: PathBuf) -> SaveData {
        SaveData {
            episode_matcher: None,
            subseries: HashMap::new(),
            path,
        }
    }

    pub fn get_mut_subseries_entry<S>(&mut self, name: S) -> &mut SubSeriesData
    where
        S: Into<String>,
    {
        self.subseries
            .entry(name.into())
            .or_insert_with(SubSeriesData::new)
    }

    pub fn from_dir(path: &Path) -> Result<SaveData, SeriesError> {
        let path = PathBuf::from(path).join(SaveData::DATA_FILE_NAME);

        if !path.exists() {
            return Ok(SaveData::new(path));
        }

        let file_contents = fs::read_to_string(&path)?;

        let mut save_data: SaveData = toml::from_str(&file_contents)?;
        save_data.path = path;

        Ok(save_data)
    }

    pub fn write_to_file(&self) -> Result<(), SeriesError> {
        let toml = toml::to_string(self)?;
        fs::write(&self.path, toml)?;

        Ok(())
    }
}

#[derive(Serialize, Deserialize)]
pub struct SubSeriesData {
    pub files_title: Option<String>,
    pub season_states: Vec<SeasonState>,
}

impl SubSeriesData {
    pub fn new() -> SubSeriesData {
        SubSeriesData {
            files_title: None,
            season_states: Vec::new(),
        }
    }
}

pub type EpisodePaths = HashMap<u32, PathBuf>;
pub type SeriesEpisodes = HashMap<String, EpisodePaths>;

pub struct SeriesInfo {
    pub name: String,
    pub episodes: HashMap<u32, PathBuf>,
}

impl SeriesInfo {
    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.episodes
            .get(&episode)
            .ok_or_else(|| SeriesError::EpisodeNotFound(episode))
    }

    pub fn select_from_subseries(
        ep_data: &mut SeriesEpisodes,
        subseries: &SubSeriesData,
    ) -> Result<Option<SeriesInfo>, SeriesError> {
        if let Some(name) = &subseries.files_title {
            let entry = ep_data.remove_entry(name);

            match entry {
                Some((name, episodes)) => return Ok(Some(SeriesInfo { name, episodes })),
                None => return Err(SeriesError::NoSeriesEpisodes),
            }
        }

        Ok(None)
    }

    pub fn prompt_select_from_episodes(info: SeriesEpisodes) -> Result<SeriesInfo, SeriesError> {
        if info.is_empty() {
            return Err(SeriesError::NoSeriesFound);
        }

        let mut info = info
            .into_iter()
            .map(|(name, eps)| SeriesInfo {
                name,
                episodes: eps,
            })
            .collect::<Vec<_>>();

        if info.len() == 1 {
            return Ok(info.swap_remove(0));
        }

        println!("multiple series found in directory");
        println!("please enter the number next to the episode files you want to use:");

        for (i, series) in info.iter().enumerate() {
            println!("{} [{}]", 1 + i, series.name);
        }

        let index = input::read_range(1, info.len())? - 1;
        let series = info.swap_remove(index);

        Ok(series)
    }
}

struct EpisodeFile {
    series_name: String,
    episode_num: u32,
}

impl EpisodeFile {
    fn parse(path: &Path, matcher: &Regex) -> Result<EpisodeFile, SeriesError> {
        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?;

        let caps = matcher
            .captures(&filename)
            .ok_or(SeriesError::EpisodeRegexCaptureFailed)?;

        let series_name = caps
            .name("name")
            .ok_or_else(|| SeriesError::UnknownRegexCapture("name".into()))
            .map(|name| EpisodeFile::cleanup_title(name.as_str()))?;

        let episode = match caps.name("episode") {
            Some(capture) => capture
                .as_str()
                .parse()
                .map_err(SeriesError::EpisodeNumParseFailed)?,
            None => 1,
        };

        Ok(EpisodeFile {
            series_name,
            episode_num: episode,
        })
    }

    fn cleanup_title(title: &str) -> String {
        let mut clean_title = String::from(title);
        let bytes = unsafe { clean_title.as_bytes_mut() };

        for byte in bytes {
            match byte {
                b'.' | b'_' => *byte = b' ',
                _ => (),
            }
        }

        clean_title.trim_right().to_string()
    }
}

// This default pattern will match episodes in several common formats, such as:
// [Group] Series Name - 01.mkv
// [Group]_Series_Name_-_01.mkv
// [Group].Series.Name.-.01.mkv
// [Group] Series Name - 01 [tag 1][tag 2].mkv
// [Group]_Series_Name_-_01_[tag1][tag2].mkv
// [Group].Series.Name.-.01.[tag1][tag2].mkv
// Series Name - 01.mkv
// Series_Name_-_01.mkv
// Series.Name.-.01.mkv
const EP_FORMAT_REGEX: &str =
    r"(?:\[.+?\](?:_+|\.+|\s*))?(?P<name>.+)(?:\s*|_*|\.*)(?:-|\.|_).*?(?P<episode>\d+)";

fn format_episode_parser_regex<'a, S>(pattern: Option<S>) -> Result<Cow<'a, Regex>, SeriesError>
where
    S: AsRef<str>,
{
    lazy_static! {
        static ref EP_FORMAT: Regex = Regex::new(EP_FORMAT_REGEX).unwrap();
    }

    match pattern {
        Some(pattern) => {
            let pattern = pattern
                .as_ref()
                .replace("{name}", "(?P<name>.+?)")
                .replace("{episode}", r"(?P<episode>\d+)");

            let regex = Regex::new(&pattern)?;
            Ok(Cow::Owned(regex))
        }
        None => Ok(Cow::Borrowed(&*EP_FORMAT)),
    }
}

pub fn parse_episode_files<S>(
    path: &Path,
    pattern: Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str>,
{
    if !path.is_dir() {
        return Err(SeriesError::NotADirectory(path.to_string_lossy().into()));
    }

    let pattern = format_episode_parser_regex(pattern)?;
    let mut data = HashMap::with_capacity(1);

    for entry in fs::read_dir(path).map_err(SeriesError::Io)? {
        let entry = entry.map_err(SeriesError::Io)?.path();

        if !entry.is_file() {
            continue;
        }

        let episode = match EpisodeFile::parse(&entry, pattern.as_ref()) {
            Ok(episode) => episode,
            Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
            Err(err) => return Err(err),
        };

        let series = data
            .entry(episode.series_name)
            .or_insert_with(|| HashMap::with_capacity(1));

        series.insert(episode.episode_num, entry);
    }

    if data.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    Ok(data)
}

pub fn parse_episode_files_until_valid<S>(
    path: &Path,
    pattern: &mut Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str> + From<String>,
{
    loop {
        match parse_episode_files(path, pattern.as_ref()) {
            Ok(data) => return Ok(data),
            Err(SeriesError::NoSeriesFound) => {
                eprintln!("error: no series found");
                println!("you will now be prompted to enter a custom regex pattern");
                println!("when entering the pattern, please mark the series name and episode number with {{name}} and {{episode}}, respectively");
                println!("note: if you're trying to play a single file (like a movie) and it doesn't have an episode number, you can omit the {{episode}} marker");
                println!("example:");
                println!("  filename: [SubGroup] Series Name - Ep01 [1080p].mkv");
                println!(r"  pattern: \[.+?\] {{name}} - Ep{{episode}} \[1080p\].mkv");
                println!("please enter your custom pattern:");
            }
            Err(err @ SeriesError::Regex(_)) | Err(err @ SeriesError::UnknownRegexCapture(_)) => {
                eprintln!("error parsing regex pattern: {}", err);
                println!("please try again:");
            }
            Err(err) => return Err(err),
        }

        let line = input::read_line()?;
        *pattern = Some(line.into());
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::fs;
use std::io::prelude::*;
use std::path::Path;
use std::fmt;

#[derive(Clone)]
#[derive(Debug)]
pub enum CellState
{
	Dead,
	Alive
}

pub struct Simulator {
	iteration_num: u32,
	height: usize,
	width: usize,
	output_dir: String,
	starting_states: Vec<Vec<CellState>>,
	current_iteration: u32,
}

impl<'a> Simulator {
	pub fn new(iteration_num: u32, output_dir: String,
		starting_states: Vec<Vec<CellState>>) -> Simulator {
		// Check input data
		if starting_states.len() < 1 {
			panic!("not enough rows");
		}
		let row_length = starting_states[0].len();
		for row in &starting_states {
			if row.len() != row_length {
				panic!("row lengths do not match");
			}
		}
		// Check input directory
		if !fs::metadata(&output_dir).is_ok() {
			match fs::create_dir_all(&output_dir) {
				Ok(_) => {},
				Err(e) => {
					panic!("output directory did not exist and cannot be created: {}",
						e.description());
				}
			}
		}
		// Assign struct
		Simulator {
			iteration_num: iteration_num,
			height: starting_states.len(),
			width: starting_states[0].len(),
			output_dir: output_dir,
			starting_states: starting_states,
			current_iteration: 0,
		}
	}

	pub fn run_simulation(&mut self) {
		// Set up first cells
		let mut current_states = Vec::new();
		for x in 0..self.starting_states.len() {
			let mut row = Vec::new();
			for y in 0..self.starting_states[x].len() {
				row.push(Cell {
					state: self.starting_states[x][y].clone(),
					row: x,
					col: y,
				});
			}
			current_states.push(row);
		}
		loop {
			// Create new states
			let mut new_states = Vec::new();
			for x in 0..self.height {
				let mut new_row = Vec::new();
				for y in 0..self.width {
					let cell = current_states[x][y].clone();
					let new_cell = cell.iterate(&current_states);
					new_row.push(new_cell);
				}
				new_states.push(new_row);
			}
			// Set current states
			current_states = new_states;
			// Output
			{
				let current_states = &current_states;
				self.output(current_states);
			}
			// Increment iteration
			self.current_iteration += 1;
			if self.current_iteration == self.iteration_num {
				break;
			}
		}
	}

	fn output(&mut self, current_states: &Vec<Vec<Cell>>) {
		let path = Path::new(&self.output_dir);
		let path = path.join(format!("{}.csv", self.current_iteration));
		let mut file = match fs::File::create(&path) {
			Err(why) => panic!("couldn't create {}: {}",
								path.display(), why.description()),
			Ok(file) => file,
		};
		for row in current_states {
			let mut row_str: String = row.iter().map(|c| match c.state {
				CellState::Dead => "0",
				CellState::Alive => "1",
			}).collect::<Vec<&str>>().join(",");
			row_str.push('\n');
			match file.write_all(row_str.as_bytes()) {
				Err(why) => {
					panic!("couldn't write to {}: {}", path.display(),
						   why.description())
				},
				Ok(_) => {},
			}
		}
	}
}

#[derive(Clone)]
struct Cell {
	state: CellState,
	row: usize,
	col: usize,
}

impl Cell {
	fn iterate(self, prev_state: &Vec<Vec<Cell>>) -> Cell {
		// TODO: write rules for Conways Game of Life
		return Cell {
			state: match self.state {
				CellState::Dead => CellState::Alive,
				CellState::Alive => CellState::Dead,
			},
			row: self.row,
			col: self.col,
		};
	}
}<commit_msg>Refectored initial state creation into separate function<commit_after>use std::error::Error;
use std::fs;
use std::io::prelude::*;
use std::path::Path;

#[derive(Clone)]
#[derive(Debug)]
pub enum CellState
{
	Dead,
	Alive
}

pub struct Simulator {
	iteration_num: u32,
	height: usize,
	width: usize,
	output_dir: String,
	starting_states: Vec<Vec<CellState>>,
	current_iteration: u32,
}

impl<'a> Simulator {
	pub fn new(iteration_num: u32, output_dir: String,
		starting_states: Vec<Vec<CellState>>) -> Simulator {
		// Check input data
		if starting_states.len() < 1 {
			panic!("not enough rows");
		}
		let row_length = starting_states[0].len();
		for row in &starting_states {
			if row.len() != row_length {
				panic!("row lengths do not match");
			}
		}
		// Check input directory
		if !fs::metadata(&output_dir).is_ok() {
			match fs::create_dir_all(&output_dir) {
				Ok(_) => {},
				Err(e) => {
					panic!("output directory did not exist and cannot be created: {}",
						e.description());
				}
			}
		}
		// Assign struct
		Simulator {
			iteration_num: iteration_num,
			height: starting_states.len(),
			width: starting_states[0].len(),
			output_dir: output_dir,
			starting_states: starting_states,
			current_iteration: 0,
		}
	}

	pub fn run_simulation(&mut self) {
		// Set up first cells
		let mut current_states = self.create_initial_states();
		loop {
			// Create new states
			let mut new_states = Vec::new();
			for x in 0..self.height {
				let mut new_row = Vec::new();
				for y in 0..self.width {
					let cell = current_states[x][y].clone();
					let new_cell = cell.iterate(&current_states);
					new_row.push(new_cell);
				}
				new_states.push(new_row);
			}
			// Set current states
			current_states = new_states;
			// Output
			{
				let current_states = &current_states;
				self.output(current_states);
			}
			// Increment iteration
			self.current_iteration += 1;
			if self.current_iteration == self.iteration_num {
				break;
			}
		}
	}

	fn create_initial_states(&mut self) -> Vec<Vec<Cell>> {
		let mut states = Vec::new();
		for x in 0..self.starting_states.len() {
			let mut row = Vec::new();
			for y in 0..self.starting_states[x].len() {
				row.push(Cell {
					state: self.starting_states[x][y].clone(),
					row: x,
					col: y,
				});
			}
			states.push(row);
		}
		return states;
	}

	fn output(&mut self, current_states: &Vec<Vec<Cell>>) {
		let path = Path::new(&self.output_dir);
		let path = path.join(format!("{}.csv", self.current_iteration));
		let mut file = match fs::File::create(&path) {
			Err(why) => panic!("couldn't create {}: {}",
								path.display(), why.description()),
			Ok(file) => file,
		};
		for row in current_states {
			let mut row_str: String = row.iter().map(|c| match c.state {
				CellState::Dead => "0",
				CellState::Alive => "1",
			}).collect::<Vec<&str>>().join(",");
			row_str.push('\n');
			match file.write_all(row_str.as_bytes()) {
				Err(why) => {
					panic!("couldn't write to {}: {}", path.display(),
						   why.description())
				},
				Ok(_) => {},
			}
		}
	}
}

#[derive(Clone)]
struct Cell {
	state: CellState,
	row: usize,
	col: usize,
}

impl Cell {
	fn iterate(self, prev_state: &Vec<Vec<Cell>>) -> Cell {
		// TODO: write rules for Conways Game of Life
		return Cell {
			state: match self.state {
				CellState::Dead => CellState::Alive,
				CellState::Alive => CellState::Dead,
			},
			row: self.row,
			col: self.col,
		};
	}
}<|endoftext|>


<commit_before>use futures::stream::Stream;
use rdkafka::client::EmptyContext;
use rdkafka::config::{ClientConfig, TopicConfig};
use rdkafka::consumer::stream_consumer::{MessageStream, StreamConsumer};
use rdkafka::consumer::{Consumer, EmptyConsumerContext};
use rdkafka::producer::{FutureProducer, FutureProducerTopic};
use rdkafka::error::KafkaError;
use rdkafka::message::Message;
use serde::de::Deserialize;
use serde::ser::Serialize;
use serde_cbor;

use std::collections::{HashMap, HashSet};
use std::hash::Hash;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;

use error::*;
use utils::format_error_chain;
use metadata::{BrokerId, ClusterId, Metadata};
use metrics::BrokerMetrics;

#[derive(Serialize, Deserialize, Debug, Hash, Eq, PartialEq)]
struct WrappedKey(String, Vec<u8>);

impl WrappedKey {
    fn new<K>(cache_name: String, key: &K) -> WrappedKey
            where K: Serialize + Deserialize {
        WrappedKey(cache_name, serde_cbor::to_vec(key).unwrap())  //TODO: error handling
    }

    pub fn cache_name(&self) -> &str {
        &self.0
    }

    pub fn serialized_key(&self) -> &[u8] {
        &self.1
    }
}

//
// ********* REPLICA WRITER **********
//

pub struct ReplicaWriter {
//    brokers: String,
//    topic_name: String,
    producer_topic: FutureProducerTopic<EmptyContext>,
}

impl ReplicaWriter {
    pub fn new(brokers: &str, topic_name: &str) -> Result<ReplicaWriter> {
        let producer = ClientConfig::new()
            .set("bootstrap.servers", brokers)
            .set("compression.codec", "gzip")
            .set("message.max.bytes", "10000000")
            .create::<FutureProducer<_>>()
            .expect("Producer creation error");

        producer.start();

        let topic = producer.get_topic(topic_name, &TopicConfig::new())
            .expect("Topic creation error");

        let writer = ReplicaWriter {
//            brokers: brokers.to_owned(),
//            topic_name: topic_name.to_owned(),
            producer_topic: topic,
        };

        Ok(writer)
    }

    // TODO: use structure for value
    pub fn write_update<K, V>(&self, name: &str, key: &K, value: &V) -> Result<()>
            where K: Serialize + Deserialize + Clone,
                  V: Serialize + Deserialize {
        let serialized_key = serde_cbor::to_vec(&WrappedKey::new(name.to_owned(), key))
            .chain_err(|| "Failed to serialize key")?;
        let serialized_value = serde_cbor::to_vec(&value)
            .chain_err(|| "Failed to serialize value")?;
        trace!("Serialized value size: {}", serialized_value.len());
        let _f = self.producer_topic.send_copy(None, Some(&serialized_value), Some(&serialized_key))
            .chain_err(|| "Failed to produce message")?;
        // _f.wait();  // Uncomment to make production synchronous
        Ok(())
    }
}

//
// ********* REPLICA READER **********
//

pub trait UpdateReceiver: Send + 'static {
    fn update(&self, name: &str, key_bytes: &[u8], msg: Message) -> Result<()>;
}

type ReplicaConsumer = StreamConsumer<EmptyConsumerContext>;

pub struct ReplicaReader {
    consumer: ReplicaConsumer,
    thread_handle: Option<thread::JoinHandle<()>>,
}

impl ReplicaReader {
    pub fn new(brokers: &str, topic_name: &str) -> Result<ReplicaReader> {
        let mut consumer: ReplicaConsumer = ClientConfig::new()
            .set("group.id", "kafka_web_replica_reader")  // TODO: make random
            .set("bootstrap.servers", brokers)
            .set("session.timeout.ms", "6000")
            .set("enable.auto.commit", "false")
            .set_default_topic_config(
                TopicConfig::new()
                .set("auto.offset.reset", "smallest")
                .finalize())
            .create()
            .chain_err(|| "Consumer creation failed")?;

        //let topic_partition = TopicPartitionList::with_topics(&vec![topic_name]);
        // TODO: does this handle partition creation?
        // consumer.assign(&topic_partition)
        consumer.subscribe(&vec![topic_name])
            .chain_err(|| "Can't subscribe to specified topics")?;

        Ok(ReplicaReader {
            consumer: consumer,
            thread_handle: None,
        })
    }

    pub fn start<R: UpdateReceiver>(&mut self, rec: R) -> Result<()> {
        let stream = self.consumer.start();
        let handle = thread::Builder::new()
            .name("replica consumer".to_string())
            .spawn(move || {
                match last_message_per_key(stream) {
                    Err(e) => format_error_chain(e),
                    Ok(state) => {
                        for (w_key, message) in state {
                            if let Err(e) = rec.update(w_key.cache_name(), w_key.serialized_key(), message) {
                                format_error_chain(e);
                            }
                        }
                    }
                };
            })
            .chain_err(|| "Failed to start polling thread")?;
        self.thread_handle = Some(handle);
        Ok(())
    }
}

fn last_message_per_key(stream: MessageStream) -> Result<HashMap<WrappedKey, Message>> {
    let mut EOF_set = HashSet::new();
    let mut state: HashMap<WrappedKey, Message> = HashMap::new();

    trace!("Started creating state");
    for message in stream.wait() {
        match message {
            Ok(Ok(m)) => {
                match parse_message_key(&m).chain_err(|| "Failed to parse message key") {
                    Ok(wrapped_key) => { state.insert(wrapped_key, m); () },
                    Err(e) => format_error_chain(e),
                };
            },
            Ok(Err(KafkaError::PartitionEOF(p))) => { EOF_set.insert(p); () },
            Ok(Err(e)) => error!("Error while reading from Kafka: {}", e),
            Err(_) => error!("Stream receive error"),
        };
        if EOF_set.len() == 2 { // TODO: make configurable
            break; // TODO: should stop consumer
        }
    }
    trace!("State creation terminated");

    Ok(state)
}

fn parse_message_key(message: &Message) -> Result<WrappedKey> {
    let key_bytes = match message.key() {
        Some(k) => k,
        None => bail!("Empty key found"),
    };

    let wrapped_key = serde_cbor::from_slice::<WrappedKey>(key_bytes)
        .chain_err(|| "Failed to decode wrapped key")?;
    Ok(wrapped_key)
}

// pub trait ReplicatedCache {
//     type Key: Serialize + Deserialize;
//     type Value: Serialize + Deserialize;
//
//     fn new(SharedReplicaWriter, &str) -> Self;
//     fn name(&self) -> &str;
//     fn insert(&self, Self::Key, Self::Value) -> Result<Arc<Self::Value>>;
//     fn get(&self, &Self::Key) -> Option<Arc<Self::Value>>;
//     fn keys(&self) -> Vec<Self::Key>;
// }

// impl<V> ValueContainer<V>
//   where K: Serialize + Deserialize {
//     fn new(id: i32, value: V) -> WrappedKey<V> {
//         ValueContainer {
//             id: id,
//             value: value,
//         }
//     }
// }

//
// ********** REPLICATEDMAP **********
//

pub struct ReplicatedMap<K, V>
        where K: Eq + Hash + Clone + Serialize + Deserialize,
              V: Clone + Serialize + Deserialize {
    name: String,
    cache_lock: Arc<RwLock<HashMap<K, RwLock<V>>>>,
    replica_writer: Arc<ReplicaWriter>,
}

impl<K, V> ReplicatedMap<K, V> where K: Eq + Hash + Clone + Serialize + Deserialize,
                                     V: Clone + Serialize + Deserialize {
    pub fn new(name: &str, replica_writer: Arc<ReplicaWriter>) -> ReplicatedMap<K, V> {
        ReplicatedMap {
            name: name.to_owned(),
            cache_lock: Arc::new(RwLock::new(HashMap::new())),
            replica_writer: replica_writer,
        }
    }

    pub fn alias(&self) -> ReplicatedMap<K, V> {
        ReplicatedMap {
            name: self.name.clone(),
            cache_lock: self.cache_lock.clone(),
            replica_writer: self.replica_writer.clone(),
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn keys(&self) -> Vec<K> {
        match self.cache_lock.read() {
            Ok(ref cache) => (*cache).keys().cloned().collect::<Vec<_>>(),
            Err(_) => panic!("Poison error"),
        }
    }

    pub fn sync_value_update(&self, key: K, value: V) {
        match self.cache_lock.read() {
            Ok(cache) => {
                let cell_lock = (*cache).get(&key);
                if cell_lock.is_some() {
                    match cell_lock.unwrap().write() {
                        Ok(mut cell) => (*cell) = value,
                        Err(_) => panic!("Poison error"),
                    };
                    return;
                }
            },
            Err(_) => panic!("Poison error"),
        };
        match self.cache_lock.write() {
            Ok(mut cache) => (*cache).insert(key, RwLock::new(value)),
            Err(_) => panic!("Poison error"),
        };
    }

    pub fn insert(&self, key: K, value: V) -> Result<()> {
        self.replica_writer.write_update(&self.name, &key, &value)
            .chain_err(|| "Failed to write cache update")?;
        self.sync_value_update(key, value);
        Ok(())
    }

    pub fn get(&self, key: &K) -> Option<V> {
        match self.cache_lock.read() {
            Ok(cache) => match (*cache).get(key) {
                Some(cell_lock) => {
                    match cell_lock.read() {
                        Ok(value) => Some(value.clone()),
                        Err(_) => panic!("Poison error"),
                    }
                },
                None => None,
            },
            Err(_) => panic!("Poison error"),
        }
    }
}

//
// ********** CACHE **********
//

pub type MetadataCache = ReplicatedMap<ClusterId, Arc<Metadata>>;
pub type MetricsCache = ReplicatedMap<(ClusterId, BrokerId), BrokerMetrics>;

pub struct Cache {
    pub metadata: MetadataCache,
    pub metrics: MetricsCache,
}

impl Cache {
    pub fn new(replica_writer: ReplicaWriter) -> Cache {
        let replica_writer_arc = Arc::new(replica_writer);
        Cache {
            metadata: ReplicatedMap::new("metadata", replica_writer_arc.clone()),
            metrics: ReplicatedMap::new("metrics", replica_writer_arc),
        }
    }

    pub fn alias(&self) -> Cache {
        Cache {
            metadata: self.metadata.alias(),
            metrics: self.metrics.alias(),
        }
    }
}

impl UpdateReceiver for Cache {
    fn update(&self, cache_name: &str, key_bytes: &[u8], msg: Message) -> Result<()> {
        match cache_name.as_ref() {
            "metadata" => {
                let cluster_id = serde_cbor::from_slice::<ClusterId>(&key_bytes)
                    .chain_err(|| "Failed to parse cluster_id in key")?;
                match msg.payload() {
                    Some(bytes) => {
                        let metadata = serde_cbor::from_slice::<Metadata>(bytes)
                            .chain_err(|| "Failed to parse metadata payload")?;
                        debug!("Sync metadata cache, cluster_id: {}", cluster_id);
                        self.metadata.sync_value_update(cluster_id, Arc::new(metadata));
                    },
                    None => bail!("Delete not implemented!"),
                };
            },
            "metrics" => {
                let broker = serde_cbor::from_slice::<(ClusterId, BrokerId)>(&key_bytes)
                    .chain_err(|| "Failed to parse key")?;
                match msg.payload() {
                    Some(bytes) => {
                        let broker_metrics = serde_cbor::from_slice::<BrokerMetrics>(bytes)
                            .chain_err(|| "Failed to parse metrics payload")?;
                        debug!("Sync metadata cache, key: {:?}", broker);
                        self.metrics.sync_value_update(broker, broker_metrics);
                    },
                    None => bail!("Delete not implemented!"),
                };
            },
            _ => {
                bail!("Unknown cache name: {}", cache_name);
            },
        };
        Ok(())
    }
}

// pub struct Cache<K, V>
//   where K: Eq + Hash + Serialize + Deserialize,
//         V: Serialize + Deserialize {
//     cache_lock: Arc<RwLock<HashMap<K, V>>>,
//     on_insert: Option<Box<Fn(&K, &V)>>,
//     on_delete: Option<Box<Fn(&K, &V)>>
// }
//
// impl<K, V> Cache<K, V>
//   where K: Eq + Hash + Serialize + Deserialize,
//         V: Serialize + Deserialize {
//
//     pub fn new() -> Cache<K, V> {
//         Cache {
//             cache_lock: Arc::new(RwLock::new(HashMap::new())),
//             on_insert: None,
//             on_delete: None,
//         }
//     }
//
//     pub fn set_on_insert<'a, CB: 'static + Fn(&K, &V)>(&'a mut self, cb: CB) -> &'a mut Cache<K, V> {
//         self.on_insert = Some(Box::new(cb));
//         self
//     }
//
//     pub fn insert(&self, key: K, value: V) {
//         self.on_insert.as_ref().map(|f| (f)(&key, &value));
//         match self.cache_lock.write() {
//             Ok(mut cache_ref) => (*cache_ref).insert(key, value),
//             Err(_) => panic!("Poison error!"),
//         };
//     }
// }
<commit_msg>Simplify cache<commit_after>use futures::stream::Stream;
use rdkafka::client::EmptyContext;
use rdkafka::config::{ClientConfig, TopicConfig};
use rdkafka::consumer::stream_consumer::{MessageStream, StreamConsumer};
use rdkafka::consumer::{Consumer, EmptyConsumerContext};
use rdkafka::producer::{FutureProducer, FutureProducerTopic};
use rdkafka::error::KafkaError;
use rdkafka::message::Message;
use serde::de::Deserialize;
use serde::ser::Serialize;
use serde_cbor;

use std::collections::{HashMap, HashSet};
use std::hash::Hash;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;

use error::*;
use utils::format_error_chain;
use metadata::{BrokerId, ClusterId, Metadata};
use metrics::BrokerMetrics;

#[derive(Serialize, Deserialize, Debug, Hash, Eq, PartialEq)]
struct WrappedKey(String, Vec<u8>);

impl WrappedKey {
    fn new<K>(cache_name: String, key: &K) -> WrappedKey
            where K: Serialize + Deserialize {
        WrappedKey(cache_name, serde_cbor::to_vec(key).unwrap())  //TODO: error handling
    }

    pub fn cache_name(&self) -> &str {
        &self.0
    }

    pub fn serialized_key(&self) -> &[u8] {
        &self.1
    }
}

//
// ********* REPLICA WRITER **********
//

pub struct ReplicaWriter {
//    brokers: String,
//    topic_name: String,
    producer_topic: FutureProducerTopic<EmptyContext>,
}

impl ReplicaWriter {
    pub fn new(brokers: &str, topic_name: &str) -> Result<ReplicaWriter> {
        let producer = ClientConfig::new()
            .set("bootstrap.servers", brokers)
            .set("compression.codec", "gzip")
            .set("message.max.bytes", "10000000")
            .create::<FutureProducer<_>>()
            .expect("Producer creation error");

        producer.start();

        let topic = producer.get_topic(topic_name, &TopicConfig::new())
            .expect("Topic creation error");

        let writer = ReplicaWriter {
//            brokers: brokers.to_owned(),
//            topic_name: topic_name.to_owned(),
            producer_topic: topic,
        };

        Ok(writer)
    }

    // TODO: use structure for value
    pub fn write_update<K, V>(&self, name: &str, key: &K, value: &V) -> Result<()>
            where K: Serialize + Deserialize + Clone,
                  V: Serialize + Deserialize {
        let serialized_key = serde_cbor::to_vec(&WrappedKey::new(name.to_owned(), key))
            .chain_err(|| "Failed to serialize key")?;
        let serialized_value = serde_cbor::to_vec(&value)
            .chain_err(|| "Failed to serialize value")?;
        trace!("Serialized value size: {}", serialized_value.len());
        let _f = self.producer_topic.send_copy(None, Some(&serialized_value), Some(&serialized_key))
            .chain_err(|| "Failed to produce message")?;
        // _f.wait();  // Uncomment to make production synchronous
        Ok(())
    }
}

//
// ********* REPLICA READER **********
//

pub enum ReplicaCacheUpdate<'a> {
    Set { key: &'a[u8], payload: &'a[u8] },
    Delete { key: &'a[u8] }
}

pub trait UpdateReceiver: Send + 'static {
    fn receive_update(&self, name: &str, update: ReplicaCacheUpdate) -> Result<()>;
}

type ReplicaConsumer = StreamConsumer<EmptyConsumerContext>;

pub struct ReplicaReader {
    consumer: ReplicaConsumer,
    thread_handle: Option<thread::JoinHandle<()>>,
}

impl ReplicaReader {
    pub fn new(brokers: &str, topic_name: &str) -> Result<ReplicaReader> {
        let mut consumer: ReplicaConsumer = ClientConfig::new()
            .set("group.id", "kafka_web_replica_reader")  // TODO: make random
            .set("bootstrap.servers", brokers)
            .set("session.timeout.ms", "6000")
            .set("enable.auto.commit", "false")
            .set_default_topic_config(
                TopicConfig::new()
                .set("auto.offset.reset", "smallest")
                .finalize())
            .create()
            .chain_err(|| "Consumer creation failed")?;

        //let topic_partition = TopicPartitionList::with_topics(&vec![topic_name]);
        // TODO: does this handle partition creation?
        // consumer.assign(&topic_partition)
        consumer.subscribe(&vec![topic_name])
            .chain_err(|| "Can't subscribe to specified topics")?;

        Ok(ReplicaReader {
            consumer: consumer,
            thread_handle: None,
        })
    }

    pub fn start<R: UpdateReceiver>(&mut self, receiver: R) -> Result<()> {
        let stream = self.consumer.start();
        let handle = thread::Builder::new()
            .name("replica consumer".to_string())
            .spawn(move || {
                match last_message_per_key(stream) {
                    Err(e) => format_error_chain(e),
                    Ok(state) => {
                        for (w_key, message) in state {
                            let update = match message.payload() {
                                Some(payload) => ReplicaCacheUpdate::Set {
                                    key: w_key.serialized_key(),
                                    payload: payload
                                },
                                None => ReplicaCacheUpdate::Delete {
                                    key: w_key.serialized_key()
                                },
                            };
                            if let Err(e) = receiver.receive_update(w_key.cache_name(), update) {
                                format_error_chain(e);
                            }
                        }
                    }
                };
            })
            .chain_err(|| "Failed to start polling thread")?;
        self.thread_handle = Some(handle);
        Ok(())
    }
}

fn last_message_per_key(stream: MessageStream) -> Result<HashMap<WrappedKey, Message>> {
    let mut EOF_set = HashSet::new();
    let mut state: HashMap<WrappedKey, Message> = HashMap::new();

    trace!("Started creating state");
    for message in stream.wait() {
        match message {
            Ok(Ok(m)) => {
                match parse_message_key(&m).chain_err(|| "Failed to parse message key") {
                    Ok(wrapped_key) => { state.insert(wrapped_key, m); () },
                    Err(e) => format_error_chain(e),
                };
            },
            Ok(Err(KafkaError::PartitionEOF(p))) => { EOF_set.insert(p); () },
            Ok(Err(e)) => error!("Error while reading from Kafka: {}", e),
            Err(_) => error!("Stream receive error"),
        };
        if EOF_set.len() == 2 { // TODO: make configurable
            break; // TODO: should stop consumer
        }
    }
    trace!("State creation terminated");

    Ok(state)
}

fn parse_message_key(message: &Message) -> Result<WrappedKey> {
    let key_bytes = match message.key() {
        Some(k) => k,
        None => bail!("Empty key found"),
    };

    let wrapped_key = serde_cbor::from_slice::<WrappedKey>(key_bytes)
        .chain_err(|| "Failed to decode wrapped key")?;
    Ok(wrapped_key)
}

// pub trait ReplicatedCache {
//     type Key: Serialize + Deserialize;
//     type Value: Serialize + Deserialize;
//
//     fn new(SharedReplicaWriter, &str) -> Self;
//     fn name(&self) -> &str;
//     fn insert(&self, Self::Key, Self::Value) -> Result<Arc<Self::Value>>;
//     fn get(&self, &Self::Key) -> Option<Arc<Self::Value>>;
//     fn keys(&self) -> Vec<Self::Key>;
// }

// impl<V> ValueContainer<V>
//   where K: Serialize + Deserialize {
//     fn new(id: i32, value: V) -> WrappedKey<V> {
//         ValueContainer {
//             id: id,
//             value: value,
//         }
//     }
// }

//
// ********** REPLICATEDMAP **********
//

pub struct ReplicatedMap<K, V>
        where K: Eq + Hash + Clone + Serialize + Deserialize,
              V: Clone + Serialize + Deserialize {
    name: String,
    cache_lock: Arc<RwLock<HashMap<K, V>>>,
    replica_writer: Arc<ReplicaWriter>,
}

impl<K, V> ReplicatedMap<K, V> where K: Eq + Hash + Clone + Serialize + Deserialize,
                                     V: Clone + Serialize + Deserialize {
    pub fn new(name: &str, replica_writer: Arc<ReplicaWriter>) -> ReplicatedMap<K, V> {
        ReplicatedMap {
            name: name.to_owned(),
            cache_lock: Arc::new(RwLock::new(HashMap::new())),
            replica_writer: replica_writer,
        }
    }

    pub fn alias(&self) -> ReplicatedMap<K, V> {
        ReplicatedMap {
            name: self.name.clone(),
            cache_lock: self.cache_lock.clone(),
            replica_writer: self.replica_writer.clone(),
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn keys(&self) -> Vec<K> {
        match self.cache_lock.read() {
            Ok(ref cache) => (*cache).keys().cloned().collect::<Vec<_>>(),
            Err(_) => panic!("Poison error"),
        }
    }

    pub fn receive_update(&self, update: ReplicaCacheUpdate) -> Result<()> {
        match update {
            ReplicaCacheUpdate::Set { key, payload } => {
                let key = serde_cbor::from_slice::<K>(&key)
                    .chain_err(|| "Failed to parse key")?;
                let value = serde_cbor::from_slice::<V>(payload)
                    .chain_err(|| "Failed to parse payload")?;
                self.sync_value_update(key, value);
            },
            ReplicaCacheUpdate::Delete { key } => {
                bail!("Delete not implemented");
            }
        }
        Ok(())
    }

    pub fn sync_value_update(&self, key: K, value: V) {
        match self.cache_lock.write() {
            Ok(mut cache) => (*cache).insert(key, value),
            Err(_) => panic!("Poison error"),
        };
    }

    pub fn insert(&self, key: K, value: V) -> Result<()> {
        self.replica_writer.write_update(&self.name, &key, &value)
            .chain_err(|| "Failed to write cache update")?;
        self.sync_value_update(key, value);
        Ok(())
    }

    pub fn get(&self, key: &K) -> Option<V> {
        match self.cache_lock.read() {
            Ok(cache) => { return (*cache).get(key).map(|v| v.clone()) },
            Err(_) => panic!("Poison error"),
        };
    }
}

//
// ********** CACHE **********
//

pub type MetadataCache = ReplicatedMap<ClusterId, Arc<Metadata>>;
pub type MetricsCache = ReplicatedMap<(ClusterId, BrokerId), BrokerMetrics>;

pub struct Cache {
    pub metadata: MetadataCache,
    pub metrics: MetricsCache,
}

impl Cache {
    pub fn new(replica_writer: ReplicaWriter) -> Cache {
        let replica_writer_arc = Arc::new(replica_writer);
        Cache {
            metadata: ReplicatedMap::new("metadata", replica_writer_arc.clone()),
            metrics: ReplicatedMap::new("metrics", replica_writer_arc),
        }
    }

    pub fn alias(&self) -> Cache {
        Cache {
            metadata: self.metadata.alias(),
            metrics: self.metrics.alias(),
        }
    }
}

impl UpdateReceiver for Cache {
    fn receive_update(&self, cache_name: &str, update: ReplicaCacheUpdate) -> Result<()> {
        match cache_name.as_ref() {
            "metadata" => self.metadata.receive_update(update),
            "metrics" => self.metrics.receive_update(update),
            _ => bail!("Unknown cache name: {}", cache_name),
        };
        Ok(())
    }
}

// pub struct Cache<K, V>
//   where K: Eq + Hash + Serialize + Deserialize,
//         V: Serialize + Deserialize {
//     cache_lock: Arc<RwLock<HashMap<K, V>>>,
//     on_insert: Option<Box<Fn(&K, &V)>>,
//     on_delete: Option<Box<Fn(&K, &V)>>
// }
//
// impl<K, V> Cache<K, V>
//   where K: Eq + Hash + Serialize + Deserialize,
//         V: Serialize + Deserialize {
//
//     pub fn new() -> Cache<K, V> {
//         Cache {
//             cache_lock: Arc::new(RwLock::new(HashMap::new())),
//             on_insert: None,
//             on_delete: None,
//         }
//     }
//
//     pub fn set_on_insert<'a, CB: 'static + Fn(&K, &V)>(&'a mut self, cb: CB) -> &'a mut Cache<K, V> {
//         self.on_insert = Some(Box::new(cb));
//         self
//     }
//
//     pub fn insert(&self, key: K, value: V) {
//         self.on_insert.as_ref().map(|f| (f)(&key, &value));
//         match self.cache_lock.write() {
//             Ok(mut cache_ref) => (*cache_ref).insert(key, value),
//             Err(_) => panic!("Poison error!"),
//         };
//     }
// }
<|endoftext|>


<commit_before><commit_msg>fixed ADD_VV instruction<commit_after><|endoftext|>


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

use core::nonzero::NonZero;
use dom::bindings::codegen::Bindings::TextEncoderBinding;
use dom::bindings::codegen::Bindings::TextEncoderBinding::TextEncoderMethods;
use dom::bindings::error::Fallible;
use dom::bindings::js::Root;
use dom::bindings::reflector::{Reflector, reflect_dom_object};
use dom::bindings::str::{DOMString, USVString};
use dom::globalscope::GlobalScope;
use encoding::EncoderTrap;
use encoding::Encoding;
use encoding::all::UTF_8;
use js::jsapi::{JSContext, JSObject};
use js::jsapi::{JS_GetUint8ArrayData, JS_NewUint8Array};
use libc::uint8_t;
use std::ptr;

#[dom_struct]
pub struct TextEncoder {
    reflector_: Reflector,
}

impl TextEncoder {
    fn new_inherited() -> TextEncoder {
        TextEncoder {
            reflector_: Reflector::new(),
        }
    }

    pub fn new(global: &GlobalScope) -> Root<TextEncoder> {
        reflect_dom_object(box TextEncoder::new_inherited(),
                           global,
                           TextEncoderBinding::Wrap)
    }

    // https://encoding.spec.whatwg.org/#dom-textencoder
    pub fn Constructor(global: &GlobalScope) -> Fallible<Root<TextEncoder>> {
        Ok(TextEncoder::new(global))
    }
}

impl TextEncoderMethods for TextEncoder {
    // https://encoding.spec.whatwg.org/#dom-textencoder-encoding
    fn Encoding(&self) -> DOMString {
        DOMString::from(UTF_8.name())
    }

    #[allow(unsafe_code)]
    // https://encoding.spec.whatwg.org/#dom-textencoder-encode
    unsafe fn Encode(&self, cx: *mut JSContext, input: USVString) -> NonZero<*mut JSObject> {
        let encoded = UTF_8.encode(&input.0, EncoderTrap::Strict).unwrap();
        let length = encoded.len() as u32;
        rooted!(in(cx) let js_object = JS_NewUint8Array(cx, length));
        assert!(!js_object.is_null());
        let mut is_shared = false;
        let js_object_data: *mut uint8_t = JS_GetUint8ArrayData(js_object.get(), &mut is_shared, ptr::null());
        assert!(!is_shared);
        ptr::copy_nonoverlapping(encoded.as_ptr(), js_object_data, length as usize);
        NonZero::new(js_object.get())
    }
}
<commit_msg>Auto merge of #15510 - absoludity:master, r=SimonSapin<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use core::nonzero::NonZero;
use dom::bindings::codegen::Bindings::TextEncoderBinding;
use dom::bindings::codegen::Bindings::TextEncoderBinding::TextEncoderMethods;
use dom::bindings::error::Fallible;
use dom::bindings::js::Root;
use dom::bindings::reflector::{Reflector, reflect_dom_object};
use dom::bindings::str::{DOMString, USVString};
use dom::globalscope::GlobalScope;
use encoding::EncoderTrap;
use encoding::Encoding;
use encoding::all::UTF_8;
use js::jsapi::{JSContext, JSObject};
use js::typedarray::Uint8Array;
use std::ptr;

#[dom_struct]
pub struct TextEncoder {
    reflector_: Reflector,
}

impl TextEncoder {
    fn new_inherited() -> TextEncoder {
        TextEncoder {
            reflector_: Reflector::new(),
        }
    }

    pub fn new(global: &GlobalScope) -> Root<TextEncoder> {
        reflect_dom_object(box TextEncoder::new_inherited(),
                           global,
                           TextEncoderBinding::Wrap)
    }

    // https://encoding.spec.whatwg.org/#dom-textencoder
    pub fn Constructor(global: &GlobalScope) -> Fallible<Root<TextEncoder>> {
        Ok(TextEncoder::new(global))
    }
}

impl TextEncoderMethods for TextEncoder {
    // https://encoding.spec.whatwg.org/#dom-textencoder-encoding
    fn Encoding(&self) -> DOMString {
        DOMString::from(UTF_8.name())
    }

    #[allow(unsafe_code)]
    // https://encoding.spec.whatwg.org/#dom-textencoder-encode
    unsafe fn Encode(&self, cx: *mut JSContext, input: USVString) -> NonZero<*mut JSObject> {
        let encoded = UTF_8.encode(&input.0, EncoderTrap::Strict).unwrap();

        rooted!(in(cx) let mut js_object = ptr::null_mut());
        assert!(Uint8Array::create(cx, encoded.len() as u32, Some(encoded.as_slice()), js_object.handle_mut()).is_ok());

        NonZero::new(js_object.get())
    }
}
<|endoftext|>


<commit_before><commit_msg>Fixed u8 overflow in day 12.<commit_after><|endoftext|>


<commit_before><commit_msg>make the Event type be easier to deal with<commit_after><|endoftext|>


<commit_before><commit_msg>Still working on memory stuff :(<commit_after><|endoftext|>


<commit_before><commit_msg>test(utils): remove doctests for private functions<commit_after><|endoftext|>


<commit_before>extern crate getopts;

use std::path::PathBuf;

pub fn build_opts() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optopt(
        "n",
        "",
        "Number of passphrases to generate (default 1)",
        "NUM",
    );
    opts.optopt("e", "", "Minimum entropy (default 60)", "MINENTROPY");
    opts.optopt("l", "", "NGram length (default 3, must be > 1)", "LENGTH");
    opts.optopt(
        "w",
        "",
        "Minimum word length for corpus (default 5)",
        "LENGTH",
    );
    opts.optflag("h", "help", "Display this help and exit");

    opts
}

pub fn parse_args(
    opts: &getopts::Options,
    args: &Vec<String>,
) -> Result<(Option<PathBuf>, usize, f64, usize, usize), &'static str> {
    let matches = opts
        .parse(&args[1..])
        .map_err(|_| "Failed to parse arguments.")?;

    if matches.opt_present("h") || matches.free.len() > 1 {
        return Err("Failed to parse arguments.");
    };

    let number = parse_flag(&matches, "n", 1)?;
    let min_entropy = parse_flag(&matches, "e", 60.0)?;
    let ngram_length = parse_flag(&matches, "l", 3)?;
    let min_word_length = parse_flag(&matches, "w", 5)?;

    if ngram_length < 2 {
        return Err("Ngram length must be greater than one.");
    };

    let filename = if matches.free.is_empty() || matches.free[0] == "-" {
        None
    } else {
        Some(PathBuf::from(matches.free[0].clone()))
    };

    Ok((filename, number, min_entropy, ngram_length, min_word_length))
}

pub fn print_usage(program: &str, opts: &getopts::Options) {
    let brief = format!("Usage: {} [FILE] [options]", program);
    print!("{}", opts.usage(&brief));
}

fn parse_flag<T: ::std::str::FromStr>(
    matches: &getopts::Matches,
    flag: &str,
    default: T,
) -> Result<T, &'static str> {
    matches
        .opt_str(flag)
        .map(|c| c.parse::<T>())
        .unwrap_or(Ok(default))
        .map_err(|_| "Failed to parse flag.")
}
<commit_msg>Rename parse_flag to parse_flag_or_default<commit_after>extern crate getopts;

use std::path::PathBuf;

pub fn build_opts() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optopt(
        "n",
        "",
        "Number of passphrases to generate (default 1)",
        "NUM",
    );
    opts.optopt("e", "", "Minimum entropy (default 60)", "MINENTROPY");
    opts.optopt("l", "", "NGram length (default 3, must be > 1)", "LENGTH");
    opts.optopt(
        "w",
        "",
        "Minimum word length for corpus (default 5)",
        "LENGTH",
    );
    opts.optflag("h", "help", "Display this help and exit");

    opts
}

pub fn parse_args(
    opts: &getopts::Options,
    args: &Vec<String>,
) -> Result<(Option<PathBuf>, usize, f64, usize, usize), &'static str> {
    let matches = opts
        .parse(&args[1..])
        .map_err(|_| "Failed to parse arguments.")?;

    if matches.opt_present("h") || matches.free.len() > 1 {
        return Err("Failed to parse arguments.");
    };

    let number = parse_flag_or_default(&matches, "n", 1)?;
    let min_entropy = parse_flag_or_default(&matches, "e", 60.0)?;
    let ngram_length = parse_flag_or_default(&matches, "l", 3)?;
    let min_word_length = parse_flag_or_default(&matches, "w", 5)?;

    if ngram_length < 2 {
        return Err("Ngram length must be greater than one.");
    };

    let filename = if matches.free.is_empty() || matches.free[0] == "-" {
        None
    } else {
        Some(PathBuf::from(matches.free[0].clone()))
    };

    Ok((filename, number, min_entropy, ngram_length, min_word_length))
}

pub fn print_usage(program: &str, opts: &getopts::Options) {
    let brief = format!("Usage: {} [FILE] [options]", program);
    print!("{}", opts.usage(&brief));
}

fn parse_flag_or_default<T: ::std::str::FromStr>(
    matches: &getopts::Matches,
    flag: &str,
    default: T,
) -> Result<T, &'static str> {
    matches
        .opt_str(flag)
        .map(|c| c.parse::<T>())
        .unwrap_or(Ok(default))
        .map_err(|_| "Failed to parse flag.")
}
<|endoftext|>


<commit_before>//! C API for mp4parse module.
//!
//! Parses ISO Base Media Format aka video/mp4 streams.
//!
//! # Examples
//!
//! ```rust
//! extern crate mp4parse;
//!
//! // Minimal valid mp4 containing no tracks.
//! let data = b"\0\0\0\x0cftypmp42";
//!
//! let context = mp4parse::mp4parse_new();
//! unsafe {
//!     let rv = mp4parse::mp4parse_read(context, data.as_ptr(), data.len());
//!     assert_eq!(0, rv);
//!     mp4parse::mp4parse_free(context);
//! }
//! ```

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

use std;
use std::io::Cursor;

// Symbols we need from our rust api.
use MediaContext;
use TrackType;
use read_mp4;
use Error;
use media_time_to_ms;
use track_time_to_ms;
use SampleEntry;

const TRACK_TYPE_H264: u32 = 0;
const TRACK_TYPE_AAC:  u32 = 1;

// These structs *must* match those declared in include/mp4parse.h.
#[repr(C)]
pub struct TrackInfo {
    track_type: u32,
    track_id: u32,
    duration: u64,
    media_time: i64, // wants to be u64? understand how elst adjustment works
}

#[repr(C)]
pub struct TrackAudioInfo {
    channels: u16,
    bit_depth: u16,
    sample_rate: u32,
//    profile: i32,
//    extended_profile: i32, // check types
}

#[repr(C)]
pub struct TrackVideoInfo {
    display_width: u32,
    display_height: u32,
    image_width: u16,
    image_height: u16,
}

/// Allocate an opaque rust-side parser context.
#[no_mangle]
pub extern "C" fn mp4parse_new() -> *mut MediaContext {
    let context = Box::new(MediaContext::new());
    Box::into_raw(context)
}

/// Free a rust-side parser context.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_free(context: *mut MediaContext) {
    assert!(!context.is_null());
    let _ = Box::from_raw(context);
}

/// Feed a buffer through `read_mp4()` with the given rust-side
/// parser context, returning the number of detected tracks.
///
/// This is safe to call with NULL arguments but will crash
/// if given invalid pointers, as is usual for C.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_read(context: *mut MediaContext, buffer: *const u8, size: usize) -> i32 {
    // Validate arguments from C.
    if context.is_null() || buffer.is_null() || size < 8 {
        return -1;
    }

    let mut context: &mut MediaContext = &mut *context;

    // Wrap the buffer we've been give in a slice.
    let b = std::slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread to catch any panics.
    let task = std::thread::spawn(move || {
        match read_mp4(&mut c, &mut context) {
            Ok(_) => {},
            Err(Error::UnexpectedEOF) => {},
            Err(e) => { panic!(e); },
        }
        // Make sure the track count fits in an i32 so we can use
        // negative values for failure.
        assert!(context.tracks.len() < i32::max_value() as usize);
        context.tracks.len() as i32
    });
    task.join().unwrap_or(-1)
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_info(context: *mut MediaContext, track: i32, info: *mut TrackInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    (*info).track_type = match track.track_type {
        TrackType::Video => TRACK_TYPE_H264,
        TrackType::Audio => TRACK_TYPE_AAC,
        TrackType::Unknown => return -1,
    };

    // Maybe context & track should just have a single simple is_valid() instead?
    if context.timescale.is_none() ||
        track.timescale.is_none() ||
        track.empty_duration.is_none() ||
        track.media_time.is_none() ||
        track.duration.is_none() ||
        track.track_id.is_none() {
            return -1;
        }

    let empty_duration = media_time_to_ms(track.empty_duration.unwrap(), context.timescale.unwrap());
    (*info).media_time = track_time_to_ms(track.media_time.unwrap(), track.timescale.unwrap()) as i64 - empty_duration as i64;
    (*info).duration = track_time_to_ms(track.duration.unwrap(), track.timescale.unwrap());
    (*info).track_id = track.track_id.unwrap();

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_audio_info(context: *mut MediaContext, track: i32, info: *mut TrackAudioInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Audio => {},
        _ => return -1,
    };

    let audio = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let audio = match audio {
        &SampleEntry::Audio(ref x) => x,
        _ => return -1,
    };

    (*info).channels = audio.channelcount;
    (*info).bit_depth = audio.samplesize;
    (*info).sample_rate = audio.samplerate >> 16; // 16.16 fixed point

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_video_info(context: *mut MediaContext, track: i32, info: *mut TrackVideoInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Video => {},
        _ => return -1,
    };

    let video = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let video = match video {
        &SampleEntry::Video(ref x) => x,
        _ => return -1,
    };

    if let Some(ref tkhd) = track.tkhd {
        (*info).display_width = tkhd.width >> 16; // 16.16 fixed point
        (*info).display_height = tkhd.height >> 16; // 16.16 fixed point
    } else {
        return -1
    }
    (*info).image_width = video.width;
    (*info).image_width = video.height;

    0
}

#[test]
fn new_context() {
    let context = mp4parse_new();
    assert!(!context.is_null());
    unsafe { mp4parse_free(context); }
}

#[test]
#[should_panic(expected = "assertion failed")]
fn free_null_context() {
    unsafe { mp4parse_free(std::ptr::null_mut()); }
}

#[test]
fn arg_validation() {
    let null_buffer = std::ptr::null();
    let null_context = std::ptr::null_mut();

    let context = mp4parse_new();
    assert!(!context.is_null());

    let buffer = vec![0u8; 8];

    unsafe {
        assert_eq!(-1, mp4parse_read(null_context, null_buffer, 0));
        assert_eq!(-1, mp4parse_read(context, null_buffer, 0));
    }

    for size in 0..buffer.len() {
        println!("testing buffer length {}", size);
        unsafe {
            assert_eq!(-1, mp4parse_read(context, buffer.as_ptr(), size));
        }
    }

    unsafe { mp4parse_free(context); }
}
<commit_msg>Don't treat tracks without an edit list as invalid.<commit_after>//! C API for mp4parse module.
//!
//! Parses ISO Base Media Format aka video/mp4 streams.
//!
//! # Examples
//!
//! ```rust
//! extern crate mp4parse;
//!
//! // Minimal valid mp4 containing no tracks.
//! let data = b"\0\0\0\x0cftypmp42";
//!
//! let context = mp4parse::mp4parse_new();
//! unsafe {
//!     let rv = mp4parse::mp4parse_read(context, data.as_ptr(), data.len());
//!     assert_eq!(0, rv);
//!     mp4parse::mp4parse_free(context);
//! }
//! ```

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

use std;
use std::io::Cursor;

// Symbols we need from our rust api.
use MediaContext;
use TrackType;
use read_mp4;
use Error;
use media_time_to_ms;
use track_time_to_ms;
use SampleEntry;

const TRACK_TYPE_H264: u32 = 0;
const TRACK_TYPE_AAC:  u32 = 1;

// These structs *must* match those declared in include/mp4parse.h.
#[repr(C)]
pub struct TrackInfo {
    track_type: u32,
    track_id: u32,
    duration: u64,
    media_time: i64, // wants to be u64? understand how elst adjustment works
}

#[repr(C)]
pub struct TrackAudioInfo {
    channels: u16,
    bit_depth: u16,
    sample_rate: u32,
//    profile: i32,
//    extended_profile: i32, // check types
}

#[repr(C)]
pub struct TrackVideoInfo {
    display_width: u32,
    display_height: u32,
    image_width: u16,
    image_height: u16,
}

/// Allocate an opaque rust-side parser context.
#[no_mangle]
pub extern "C" fn mp4parse_new() -> *mut MediaContext {
    let context = Box::new(MediaContext::new());
    Box::into_raw(context)
}

/// Free a rust-side parser context.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_free(context: *mut MediaContext) {
    assert!(!context.is_null());
    let _ = Box::from_raw(context);
}

/// Feed a buffer through `read_mp4()` with the given rust-side
/// parser context, returning the number of detected tracks.
///
/// This is safe to call with NULL arguments but will crash
/// if given invalid pointers, as is usual for C.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_read(context: *mut MediaContext, buffer: *const u8, size: usize) -> i32 {
    // Validate arguments from C.
    if context.is_null() || buffer.is_null() || size < 8 {
        return -1;
    }

    let mut context: &mut MediaContext = &mut *context;

    // Wrap the buffer we've been give in a slice.
    let b = std::slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread to catch any panics.
    let task = std::thread::spawn(move || {
        match read_mp4(&mut c, &mut context) {
            Ok(_) => {},
            Err(Error::UnexpectedEOF) => {},
            Err(e) => { panic!(e); },
        }
        // Make sure the track count fits in an i32 so we can use
        // negative values for failure.
        assert!(context.tracks.len() < i32::max_value() as usize);
        context.tracks.len() as i32
    });
    task.join().unwrap_or(-1)
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_info(context: *mut MediaContext, track: i32, info: *mut TrackInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    (*info).track_type = match track.track_type {
        TrackType::Video => TRACK_TYPE_H264,
        TrackType::Audio => TRACK_TYPE_AAC,
        TrackType::Unknown => return -1,
    };

    // Maybe context & track should just have a single simple is_valid() instead?
    if context.timescale.is_none() ||
        track.timescale.is_none() ||
        track.duration.is_none() ||
        track.track_id.is_none() {
            return -1;
        }

    let empty_duration = if track.empty_duration.is_some() {
        media_time_to_ms(track.empty_duration.unwrap(), context.timescale.unwrap())
    } else {
        0
    };
    (*info).media_time = if track.media_time.is_some() {
        track_time_to_ms(track.media_time.unwrap(), track.timescale.unwrap()) as i64 - empty_duration as i64
    } else {
        0
    };
    (*info).duration = track_time_to_ms(track.duration.unwrap(), track.timescale.unwrap());
    (*info).track_id = track.track_id.unwrap();

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_audio_info(context: *mut MediaContext, track: i32, info: *mut TrackAudioInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Audio => {},
        _ => return -1,
    };

    let audio = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let audio = match audio {
        &SampleEntry::Audio(ref x) => x,
        _ => return -1,
    };

    (*info).channels = audio.channelcount;
    (*info).bit_depth = audio.samplesize;
    (*info).sample_rate = audio.samplerate >> 16; // 16.16 fixed point

    0
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_video_info(context: *mut MediaContext, track: i32, info: *mut TrackVideoInfo) -> i32 {
    if context.is_null() || track < 0 || info.is_null() {
        return -1;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return -1;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Video => {},
        _ => return -1,
    };

    let video = match track.data {
        Some(ref data) => data,
        None => return -1,
    };

    let video = match video {
        &SampleEntry::Video(ref x) => x,
        _ => return -1,
    };

    if let Some(ref tkhd) = track.tkhd {
        (*info).display_width = tkhd.width >> 16; // 16.16 fixed point
        (*info).display_height = tkhd.height >> 16; // 16.16 fixed point
    } else {
        return -1
    }
    (*info).image_width = video.width;
    (*info).image_width = video.height;

    0
}

#[test]
fn new_context() {
    let context = mp4parse_new();
    assert!(!context.is_null());
    unsafe { mp4parse_free(context); }
}

#[test]
#[should_panic(expected = "assertion failed")]
fn free_null_context() {
    unsafe { mp4parse_free(std::ptr::null_mut()); }
}

#[test]
fn arg_validation() {
    let null_buffer = std::ptr::null();
    let null_context = std::ptr::null_mut();

    let context = mp4parse_new();
    assert!(!context.is_null());

    let buffer = vec![0u8; 8];

    unsafe {
        assert_eq!(-1, mp4parse_read(null_context, null_buffer, 0));
        assert_eq!(-1, mp4parse_read(context, null_buffer, 0));
    }

    for size in 0..buffer.len() {
        println!("testing buffer length {}", size);
        unsafe {
            assert_eq!(-1, mp4parse_read(context, buffer.as_ptr(), size));
        }
    }

    unsafe { mp4parse_free(context); }
}
<|endoftext|>


<commit_before>//! Graph traits for associated data and graph construction.

use Graph;
use ::{
    EdgeType,
};
use graph::IndexType;
use visit::{
    GraphBase,
    NodeCount,
    NodeIndexable,
};

/// Define associated data for nodes and edges
pub trait Data : GraphBase {
    type NodeWeight;
    type EdgeWeight;
}

impl<'a, G> Data for &'a G
    where G: Data,
{
    type NodeWeight = G::NodeWeight;
    type EdgeWeight = G::EdgeWeight;
}

pub trait Prop : GraphBase {
    /// The kind edges in the graph.
    type EdgeType: EdgeType;

    fn is_directed(&self) -> bool {
        <Self::EdgeType>::is_directed()
    }
}

impl<N, E, Ty, Ix> Prop for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    type EdgeType = Ty;
}

pub trait DataMap : Data {
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight>;
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight>;
}

impl<'a, G> DataMap for &'a G
    where G: DataMap,
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        (*self).node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        (*self).edge_weight(id)
    }
}

pub trait DataMapMut : DataMap {
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight>;
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight>;
}

pub trait Build : Data + NodeCount {
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId;
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Self::EdgeId;
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId;
}

pub trait Create : Build + DataMapMut {
    fn with_capacity(nodes: usize, edges: usize) -> Self;
}

impl<N, E, Ty, Ix> Data for Graph<N, E, Ty, Ix>
    where Ix: IndexType
{
    type NodeWeight = N;
    type EdgeWeight = E;
}

impl<N, E, Ty, Ix> DataMap for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        self.node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.edge_weight(id)
    }
}

impl<N, E, Ty, Ix> DataMapMut for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight> {
        self.node_weight_mut(id)
    }
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight> {
        self.edge_weight_mut(id)
    }
}

impl<N, E, Ty, Ix> NodeCount for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn node_count(&self) -> usize {
        self.node_count()
    }
}

impl<N, E, Ty, Ix> Build for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.add_edge(a, b, weight)
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.update_edge(a, b, weight)
    }
}

impl<N, E, Ty, Ix> Create for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

/// A graph element
pub enum Element<N, E> {
    Node(N),
    Edge(usize, usize, E),
}

pub trait FromElements : Create {
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        let mut gr = Self::with_capacity(0, 0);
        // usize -> NodeId map
        let mut map = Vec::new();
        for element in iterable {
            match element {
                Element::Node(w) => {
                    map.push(gr.add_node(w));
                }
                Element::Edge(a, b, w) => {
                    gr.add_edge(map[a], map[b], w);
                }
            }
        }
        gr
    }

}

fn from_elements_indexable<G, I>(iterable: I) -> G
    where G: Create + NodeIndexable,
          I: IntoIterator<Item=Element<G::NodeWeight, G::EdgeWeight>>,
{
    let mut gr = G::with_capacity(0, 0);
    let map = G::from_index;
    for element in iterable {
        match element {
            Element::Node(w) => {
                gr.add_node(w);
            }
            Element::Edge(a, b, w) => {
                gr.add_edge(map(a), map(b), w);
            }
        }
    }
    gr
}

impl<N, E, Ty, Ix> FromElements for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}
<commit_msg>Change .add_edge() to return Option<commit_after>//! Graph traits for associated data and graph construction.

use Graph;
use ::{
    EdgeType,
};
use graph::IndexType;
use visit::{
    GraphBase,
    NodeCount,
    NodeIndexable,
};

/// Define associated data for nodes and edges
pub trait Data : GraphBase {
    type NodeWeight;
    type EdgeWeight;
}

impl<'a, G> Data for &'a G
    where G: Data,
{
    type NodeWeight = G::NodeWeight;
    type EdgeWeight = G::EdgeWeight;
}

pub trait Prop : GraphBase {
    /// The kind edges in the graph.
    type EdgeType: EdgeType;

    fn is_directed(&self) -> bool {
        <Self::EdgeType>::is_directed()
    }
}

impl<N, E, Ty, Ix> Prop for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    type EdgeType = Ty;
}

pub trait DataMap : Data {
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight>;
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight>;
}

impl<'a, G> DataMap for &'a G
    where G: DataMap,
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        (*self).node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        (*self).edge_weight(id)
    }
}

pub trait DataMapMut : DataMap {
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight>;
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight>;
}

pub trait Build : Data + NodeCount {
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId;
    /// Add a new edge. If parallel edges (duplicate) are not allowed and
    /// the edge already exists, return `None`.
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId> {
        Some(self.update_edge(a, b, weight))
    }
    /// Add or update the edge from `a` to `b`. Return the id of the affected
    /// edge.
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId;
}

pub trait Create : Build + DataMapMut {
    fn with_capacity(nodes: usize, edges: usize) -> Self;
}

impl<N, E, Ty, Ix> Data for Graph<N, E, Ty, Ix>
    where Ix: IndexType
{
    type NodeWeight = N;
    type EdgeWeight = E;
}

impl<N, E, Ty, Ix> DataMap for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight(&self, id: Self::NodeId) -> Option<&Self::NodeWeight> {
        self.node_weight(id)
    }
    fn edge_weight(&self, id: Self::EdgeId) -> Option<&Self::EdgeWeight> {
        self.edge_weight(id)
    }
}

impl<N, E, Ty, Ix> DataMapMut for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType
{
    fn node_weight_mut(&mut self, id: Self::NodeId) -> Option<&mut Self::NodeWeight> {
        self.node_weight_mut(id)
    }
    fn edge_weight_mut(&mut self, id: Self::EdgeId) -> Option<&mut Self::EdgeWeight> {
        self.edge_weight_mut(id)
    }
}

impl<N, E, Ty, Ix> NodeCount for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn node_count(&self) -> usize {
        self.node_count()
    }
}

impl<N, E, Ty, Ix> Build for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn add_node(&mut self, weight: Self::NodeWeight) -> Self::NodeId {
        self.add_node(weight)
    }
    fn add_edge(&mut self,
                a: Self::NodeId,
                b: Self::NodeId,
                weight: Self::EdgeWeight) -> Option<Self::EdgeId>
    {
        Some(self.add_edge(a, b, weight))
    }
    fn update_edge(&mut self,
                   a: Self::NodeId,
                   b: Self::NodeId,
                   weight: Self::EdgeWeight) -> Self::EdgeId
    {
        self.update_edge(a, b, weight)
    }
}

impl<N, E, Ty, Ix> Create for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self::with_capacity(nodes, edges)
    }
}

/// A graph element
pub enum Element<N, E> {
    Node(N),
    Edge(usize, usize, E),
}

pub trait FromElements : Create {
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        let mut gr = Self::with_capacity(0, 0);
        // usize -> NodeId map
        let mut map = Vec::new();
        for element in iterable {
            match element {
                Element::Node(w) => {
                    map.push(gr.add_node(w));
                }
                Element::Edge(a, b, w) => {
                    gr.add_edge(map[a], map[b], w);
                }
            }
        }
        gr
    }

}

fn from_elements_indexable<G, I>(iterable: I) -> G
    where G: Create + NodeIndexable,
          I: IntoIterator<Item=Element<G::NodeWeight, G::EdgeWeight>>,
{
    let mut gr = G::with_capacity(0, 0);
    let map = G::from_index;
    for element in iterable {
        match element {
            Element::Node(w) => {
                gr.add_node(w);
            }
            Element::Edge(a, b, w) => {
                gr.add_edge(map(a), map(b), w);
            }
        }
    }
    gr
}

impl<N, E, Ty, Ix> FromElements for Graph<N, E, Ty, Ix>
    where Ty: EdgeType,
          Ix: IndexType,
{
    fn from_elements<I>(iterable: I) -> Self
        where Self: Sized,
              I: IntoIterator<Item=Element<Self::NodeWeight, Self::EdgeWeight>>,
    {
        from_elements_indexable(iterable)
    }
}
<|endoftext|>