<commit_before>#[cfg(target_family = "unix")]
use directories::{BaseDirs, ProjectDirs};

use once_cell::sync::Lazy;

use std::fs::{create_dir_all, File};
use std::io::ErrorKind;
#[cfg(feature = "gettext")]
use std::io::Read;
#[cfg(feature = "gettext")]
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;

#[cfg(target_family = "unix")]
use std::fs::read_dir;
#[cfg(target_family = "unix")]
use std::io::Write;

// Remove "-application" from `CARGO_PKG_NAME`
pub static APP_NAME: Lazy<String> = Lazy::new(|| {
    env!("CARGO_PKG_NAME")
        .rsplitn(2, '-')
        .last()
        .unwrap()
        .to_string()
});

#[cfg(feature = "gettext")]
fn po_path() -> PathBuf {
    PathBuf::from("..").join("po")
}

fn res_path() -> PathBuf {
    PathBuf::from("..").join("res")
}

fn target_path() -> PathBuf {
    PathBuf::from("..").join("target")
}

fn generate_resources() {
    let output_path = target_path().join("resources");
    create_dir_all(&output_path).unwrap();

    // Icons
    let input_path = res_path().join("icons").join("hicolor");

    let mut compile_res = Command::new("glib-compile-resources");
    compile_res
        .arg("--generate")
        .arg(format!("--sourcedir={}", input_path.to_str().unwrap()))
        .arg(format!(
            "--target={}",
            output_path.join("icons.gresource").to_str().unwrap(),
        ))
        .arg(input_path.join("icons.gresource.xml").to_str().unwrap());

    match compile_res.status() {
        Ok(status) => {
            if !status.success() {
                panic!(
                    "Failed to generate resources file for icons\n{:?}",
                    compile_res,
                );
            }
        }
        Err(ref error) => match error.kind() {
            ErrorKind::NotFound => panic!(
                "Can't generate icon resources: command `glib-compile-resources` not available"
            ),
            _ => panic!("Error invoking `compile_res`: {}", error),
        },
    }

    // UI
    let input_path = res_path().join("ui");

    let mut compile_res = Command::new("glib-compile-resources");
    compile_res
        .arg("--generate")
        .arg(format!("--sourcedir={}", input_path.to_str().unwrap()))
        .arg(format!(
            "--target={}",
            output_path.join("ui.gresource").to_str().unwrap(),
        ))
        .arg(input_path.join("ui.gresource.xml").to_str().unwrap());

    match compile_res.status() {
        Ok(status) => {
            if !status.success() {
                panic!(
                    "Failed to generate resources file for the UI\n{:?}",
                    compile_res,
                );
            }
        }
        Err(ref error) => match error.kind() {
            ErrorKind::NotFound => panic!(
                "Can't generate UI resources: command `glib-compile-resources` not available"
            ),
            _ => panic!("Error invoking `compile_res`: {}", error),
        },
    }
}

#[cfg(feature = "gettext")]
fn generate_translations() {
    let mut linguas_file =
        File::open(&po_path().join("LINGUAS")).expect("Couldn't find po/LINGUAS");
    let mut linguas = String::new();
    linguas_file
        .read_to_string(&mut linguas)
        .expect("Couldn't read po/LINGUAS as string");

    for lingua in linguas.lines() {
        let mo_path = target_path()
            .join("locale")
            .join(lingua)
            .join("LC_MESSAGES");
        create_dir_all(&mo_path).unwrap();

        let mut msgfmt = Command::new("msgfmt");
        msgfmt
            .arg(format!(
                "--output-file={}",
                mo_path.join("media-toc.mo").to_str().unwrap()
            ))
            .arg(format!("--directory={}", po_path().to_str().unwrap()))
            .arg(format!("{}.po", lingua));

        match msgfmt.status() {
            Ok(status) => {
                if !status.success() {
                    panic!(
                        "Failed to generate mo file for lingua {}\n{:?}",
                        lingua, msgfmt,
                    );
                }
            }
            Err(ref error) => match error.kind() {
                ErrorKind::NotFound => {
                    eprintln!("Can't generate translations: command `msgfmt` not available");
                    return;
                }
                _ => panic!("Error invoking `msgfmt`: {}", error),
            },
        }
    }
}

// FIXME: figure out macOS conventions for icons & translations
#[cfg(target_family = "unix")]
fn generate_install_script() {
    let base_dirs = BaseDirs::new().unwrap();
    // Note: `base_dirs.executable_dir()` is `None` on macOS
    if let Some(exe_dir) = base_dirs.executable_dir() {
        let project_dirs = ProjectDirs::from("org", "fengalin", &APP_NAME).unwrap();
        let app_data_dir = project_dirs.data_dir();
        let data_dir = app_data_dir.parent().unwrap();

        match File::create(&target_path().join("install")) {
            Ok(mut install_file) => {
                install_file
                    .write_all(format!("# User install script for {}\n", *APP_NAME).as_bytes())
                    .unwrap();

                install_file.write_all(b"\n# Install executable\n").unwrap();
                install_file
                    .write_all(format!("mkdir -p {:?}\n", exe_dir).as_bytes())
                    .unwrap();
                install_file
                    .write_all(
                        format!(
                            "cp {:?} {:?}\n",
                            target_path()
                                .canonicalize()
                                .unwrap()
                                .join("release")
                                .join(&*APP_NAME),
                            exe_dir.join(&*APP_NAME),
                        )
                        .as_bytes(),
                    )
                    .unwrap();

                install_file.write_all(b"\n# Install icons\n").unwrap();
                let icon_target_dir = data_dir.join("icons").join("hicolor");
                let entry_iter = read_dir(res_path().join("icons").join("hicolor")).unwrap();
                for entry in entry_iter {
                    let entry = entry.unwrap();
                    let entry_path = entry.path();
                    if entry_path.is_dir() {
                        let target_dir = icon_target_dir.join(&entry.file_name());
                        install_file
                            .write_all(format!("mkdir -p {:?}\n", target_dir).as_bytes())
                            .unwrap();

                        install_file
                            .write_all(
                                format!(
                                    "cp -r {:?} {:?}\n",
                                    entry_path.canonicalize().unwrap().join("apps"),
                                    target_dir,
                                )
                                .as_bytes(),
                            )
                            .unwrap();
                    }
                }

                #[cfg(feature = "gettext")]
                generate_install_translations_script(&mut install_file, &data_dir);

                install_file
                    .write_all(b"\n# Install desktop file\n")
                    .unwrap();
                let desktop_target_dir = data_dir.join("applications");
                install_file
                    .write_all(format!("mkdir -p {:?}\n", desktop_target_dir).as_bytes())
                    .unwrap();
                install_file
                    .write_all(
                        format!(
                            "cp {:?} {:?}\n",
                            res_path()
                                .join(&format!("org.fengalin.{}.desktop", *APP_NAME))
                                .canonicalize()
                                .unwrap(),
                            desktop_target_dir,
                        )
                        .as_bytes(),
                    )
                    .unwrap();
            }
            Err(err) => panic!("Couldn't create file `target/install`: {:?}", err),
        }
    }
}

#[cfg(feature = "gettext")]
fn generate_install_translations_script(install_file: &mut File, data_dir: &Path) {
    install_file
        .write_all(b"\n# Install translations\n")
        .unwrap();
    install_file
        .write_all(format!("mkdir -p {:?}\n", data_dir).as_bytes())
        .unwrap();
    install_file
        .write_all(
            format!(
                "cp -r {:?} {:?}\n",
                target_path().join("locale").canonicalize().unwrap(),
                data_dir,
            )
            .as_bytes(),
        )
        .unwrap();
}

// FIXME: figure out macOS conventions for icons & translations
#[cfg(target_family = "unix")]
fn generate_uninstall_script() {
    let base_dirs = BaseDirs::new().unwrap();
    // Note: `base_dirs.executable_dir()` is `None` on macOS
    if let Some(exe_dir) = base_dirs.executable_dir() {
        let project_dirs = ProjectDirs::from("org", "fengalin", &APP_NAME).unwrap();
        let app_data_dir = project_dirs.data_dir();
        let data_dir = app_data_dir.parent().unwrap();

        match File::create(&target_path().join("uninstall")) {
            Ok(mut install_file) => {
                install_file
                    .write_all(format!("# User uninstall script for {}\n", *APP_NAME).as_bytes())
                    .unwrap();

                install_file
                    .write_all(b"\n# Uninstall executable\n")
                    .unwrap();
                install_file
                    .write_all(format!("rm {:?}\n", exe_dir.join(&*APP_NAME)).as_bytes())
                    .unwrap();
                install_file
                    .write_all(format!("rmdir -p {:?}\n", exe_dir).as_bytes())
                    .unwrap();

                install_file.write_all(b"\n# Uninstall icons\n").unwrap();
                let icon_target_dir = data_dir.join("icons").join("hicolor");
                let entry_iter = read_dir(res_path().join("icons").join("hicolor")).unwrap();
                for entry in entry_iter {
                    let entry = entry.unwrap();
                    let entry_path = entry.path();
                    if entry_path.is_dir() {
                        let icon_target_dir = icon_target_dir
                            .join(entry_path.file_name().unwrap())
                            .join("apps");
                        let icons_subdir_iter = read_dir(entry_path.join("apps")).unwrap();
                        for icon_subdir_entry in icons_subdir_iter {
                            let icon_subdir_entry = icon_subdir_entry.unwrap();
                            install_file
                                .write_all(
                                    format!(
                                        "rm {:?}\n",
                                        icon_target_dir.join(icon_subdir_entry.file_name()),
                                    )
                                    .as_bytes(),
                                )
                                .unwrap();
                        }

                        install_file
                            .write_all(format!("rmdir -p {:?}\n", icon_target_dir).as_bytes())
                            .unwrap();
                    }
                }

                #[cfg(feature = "gettext")]
                generate_uninstall_translations_script(&mut install_file, &data_dir);

                install_file
                    .write_all(b"\n# Uninstall desktop file\n")
                    .unwrap();
                let desktop_target_dir = data_dir.join("applications");
                install_file
                    .write_all(
                        format!(
                            "rm {:?}\n",
                            desktop_target_dir.join(&format!("org.fengalin.{}.desktop", *APP_NAME)),
                        )
                        .as_bytes(),
                    )
                    .unwrap();
                install_file
                    .write_all(format!("rmdir -p {:?}\n", desktop_target_dir).as_bytes())
                    .unwrap();
            }
            Err(err) => panic!("Couldn't create file `target/uninstall`: {:?}", err),
        }
    }
}

#[cfg(feature = "gettext")]
fn generate_uninstall_translations_script(install_file: &mut File, data_dir: &Path) {
    if let Ok(mut linguas_file) = File::open(&po_path().join("LINGUAS")) {
        let mut linguas = String::new();
        linguas_file
            .read_to_string(&mut linguas)
            .expect("Couldn't read po/LINGUAS as string");

        install_file
            .write_all(b"\n# Uninstall translations\n")
            .unwrap();
        let locale_base_dir = data_dir.join("locale");
        for lingua in linguas.lines() {
            let lingua_dir = locale_base_dir.join(lingua).join("LC_MESSAGES");
            install_file
                .write_all(
                    format!("rm {:?}\n", lingua_dir.join(&format!("{}.mo", *APP_NAME)),).as_bytes(),
                )
                .unwrap();

            install_file
                .write_all(format!("rmdir -p {:?}\n", lingua_dir).as_bytes())
                .unwrap();
        }
    }
}

fn main() {
    generate_resources();
    #[cfg(feature = "gettext")]
    generate_translations();

    #[cfg(target_family = "unix")]
    generate_install_script();

    #[cfg(target_family = "unix")]
    generate_uninstall_script();
}
<commit_msg>application: fix clippy warning in build.rs<commit_after>#[cfg(target_family = "unix")]
use directories::{BaseDirs, ProjectDirs};

use once_cell::sync::Lazy;

use std::fs::{create_dir_all, File};
use std::io::ErrorKind;
#[cfg(feature = "gettext")]
use std::io::Read;
#[cfg(feature = "gettext")]
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;

#[cfg(target_family = "unix")]
use std::fs::read_dir;
#[cfg(target_family = "unix")]
use std::io::Write;

// Remove "-application" from `CARGO_PKG_NAME`
pub static APP_NAME: Lazy<String> = Lazy::new(|| {
    env!("CARGO_PKG_NAME")
        .rsplitn(2, '-')
        .last()
        .unwrap()
        .to_string()
});

#[cfg(feature = "gettext")]
fn po_path() -> PathBuf {
    PathBuf::from("..").join("po")
}

fn res_path() -> PathBuf {
    PathBuf::from("..").join("res")
}

fn target_path() -> PathBuf {
    PathBuf::from("..").join("target")
}

fn generate_resources() {
    let output_path = target_path().join("resources");
    create_dir_all(&output_path).unwrap();

    // Icons
    let input_path = res_path().join("icons").join("hicolor");

    let mut compile_res = Command::new("glib-compile-resources");
    compile_res
        .arg("--generate")
        .arg(format!("--sourcedir={}", input_path.to_str().unwrap()))
        .arg(format!(
            "--target={}",
            output_path.join("icons.gresource").to_str().unwrap(),
        ))
        .arg(input_path.join("icons.gresource.xml").to_str().unwrap());

    match compile_res.status() {
        Ok(status) => {
            if !status.success() {
                panic!(
                    "Failed to generate resources file for icons\n{:?}",
                    compile_res,
                );
            }
        }
        Err(ref error) => match error.kind() {
            ErrorKind::NotFound => panic!(
                "Can't generate icon resources: command `glib-compile-resources` not available"
            ),
            _ => panic!("Error invoking `compile_res`: {}", error),
        },
    }

    // UI
    let input_path = res_path().join("ui");

    let mut compile_res = Command::new("glib-compile-resources");
    compile_res
        .arg("--generate")
        .arg(format!("--sourcedir={}", input_path.to_str().unwrap()))
        .arg(format!(
            "--target={}",
            output_path.join("ui.gresource").to_str().unwrap(),
        ))
        .arg(input_path.join("ui.gresource.xml").to_str().unwrap());

    match compile_res.status() {
        Ok(status) => {
            if !status.success() {
                panic!(
                    "Failed to generate resources file for the UI\n{:?}",
                    compile_res,
                );
            }
        }
        Err(ref error) => match error.kind() {
            ErrorKind::NotFound => panic!(
                "Can't generate UI resources: command `glib-compile-resources` not available"
            ),
            _ => panic!("Error invoking `compile_res`: {}", error),
        },
    }
}

#[cfg(feature = "gettext")]
fn generate_translations() {
    let mut linguas_file =
        File::open(&po_path().join("LINGUAS")).expect("Couldn't find po/LINGUAS");
    let mut linguas = String::new();
    linguas_file
        .read_to_string(&mut linguas)
        .expect("Couldn't read po/LINGUAS as string");

    for lingua in linguas.lines() {
        let mo_path = target_path()
            .join("locale")
            .join(lingua)
            .join("LC_MESSAGES");
        create_dir_all(&mo_path).unwrap();

        let mut msgfmt = Command::new("msgfmt");
        msgfmt
            .arg(format!(
                "--output-file={}",
                mo_path.join("media-toc.mo").to_str().unwrap()
            ))
            .arg(format!("--directory={}", po_path().to_str().unwrap()))
            .arg(format!("{}.po", lingua));

        match msgfmt.status() {
            Ok(status) => {
                if !status.success() {
                    panic!(
                        "Failed to generate mo file for lingua {}\n{:?}",
                        lingua, msgfmt,
                    );
                }
            }
            Err(ref error) => match error.kind() {
                ErrorKind::NotFound => {
                    eprintln!("Can't generate translations: command `msgfmt` not available");
                    return;
                }
                _ => panic!("Error invoking `msgfmt`: {}", error),
            },
        }
    }
}

// FIXME: figure out macOS conventions for icons & translations
#[cfg(target_family = "unix")]
fn generate_install_script() {
    let base_dirs = BaseDirs::new().unwrap();
    // Note: `base_dirs.executable_dir()` is `None` on macOS
    if let Some(exe_dir) = base_dirs.executable_dir() {
        let project_dirs = ProjectDirs::from("org", "fengalin", &APP_NAME).unwrap();
        let app_data_dir = project_dirs.data_dir();
        let data_dir = app_data_dir.parent().unwrap();

        match File::create(&target_path().join("install")) {
            Ok(mut install_file) => {
                install_file
                    .write_all(format!("# User install script for {}\n", *APP_NAME).as_bytes())
                    .unwrap();

                install_file.write_all(b"\n# Install executable\n").unwrap();
                install_file
                    .write_all(format!("mkdir -p {:?}\n", exe_dir).as_bytes())
                    .unwrap();
                install_file
                    .write_all(
                        format!(
                            "cp {:?} {:?}\n",
                            target_path()
                                .canonicalize()
                                .unwrap()
                                .join("release")
                                .join(&*APP_NAME),
                            exe_dir.join(&*APP_NAME),
                        )
                        .as_bytes(),
                    )
                    .unwrap();

                install_file.write_all(b"\n# Install icons\n").unwrap();
                let icon_target_dir = data_dir.join("icons").join("hicolor");
                let entry_iter = read_dir(res_path().join("icons").join("hicolor")).unwrap();
                for entry in entry_iter {
                    let entry = entry.unwrap();
                    let entry_path = entry.path();
                    if entry_path.is_dir() {
                        let target_dir = icon_target_dir.join(&entry.file_name());
                        install_file
                            .write_all(format!("mkdir -p {:?}\n", target_dir).as_bytes())
                            .unwrap();

                        install_file
                            .write_all(
                                format!(
                                    "cp -r {:?} {:?}\n",
                                    entry_path.canonicalize().unwrap().join("apps"),
                                    target_dir,
                                )
                                .as_bytes(),
                            )
                            .unwrap();
                    }
                }

                #[cfg(feature = "gettext")]
                generate_install_translations_script(&mut install_file, data_dir);

                install_file
                    .write_all(b"\n# Install desktop file\n")
                    .unwrap();
                let desktop_target_dir = data_dir.join("applications");
                install_file
                    .write_all(format!("mkdir -p {:?}\n", desktop_target_dir).as_bytes())
                    .unwrap();
                install_file
                    .write_all(
                        format!(
                            "cp {:?} {:?}\n",
                            res_path()
                                .join(&format!("org.fengalin.{}.desktop", *APP_NAME))
                                .canonicalize()
                                .unwrap(),
                            desktop_target_dir,
                        )
                        .as_bytes(),
                    )
                    .unwrap();
            }
            Err(err) => panic!("Couldn't create file `target/install`: {:?}", err),
        }
    }
}

#[cfg(feature = "gettext")]
fn generate_install_translations_script(install_file: &mut File, data_dir: &Path) {
    install_file
        .write_all(b"\n# Install translations\n")
        .unwrap();
    install_file
        .write_all(format!("mkdir -p {:?}\n", data_dir).as_bytes())
        .unwrap();
    install_file
        .write_all(
            format!(
                "cp -r {:?} {:?}\n",
                target_path().join("locale").canonicalize().unwrap(),
                data_dir,
            )
            .as_bytes(),
        )
        .unwrap();
}

// FIXME: figure out macOS conventions for icons & translations
#[cfg(target_family = "unix")]
fn generate_uninstall_script() {
    let base_dirs = BaseDirs::new().unwrap();
    // Note: `base_dirs.executable_dir()` is `None` on macOS
    if let Some(exe_dir) = base_dirs.executable_dir() {
        let project_dirs = ProjectDirs::from("org", "fengalin", &APP_NAME).unwrap();
        let app_data_dir = project_dirs.data_dir();
        let data_dir = app_data_dir.parent().unwrap();

        match File::create(&target_path().join("uninstall")) {
            Ok(mut install_file) => {
                install_file
                    .write_all(format!("# User uninstall script for {}\n", *APP_NAME).as_bytes())
                    .unwrap();

                install_file
                    .write_all(b"\n# Uninstall executable\n")
                    .unwrap();
                install_file
                    .write_all(format!("rm {:?}\n", exe_dir.join(&*APP_NAME)).as_bytes())
                    .unwrap();
                install_file
                    .write_all(format!("rmdir -p {:?}\n", exe_dir).as_bytes())
                    .unwrap();

                install_file.write_all(b"\n# Uninstall icons\n").unwrap();
                let icon_target_dir = data_dir.join("icons").join("hicolor");
                let entry_iter = read_dir(res_path().join("icons").join("hicolor")).unwrap();
                for entry in entry_iter {
                    let entry = entry.unwrap();
                    let entry_path = entry.path();
                    if entry_path.is_dir() {
                        let icon_target_dir = icon_target_dir
                            .join(entry_path.file_name().unwrap())
                            .join("apps");
                        let icons_subdir_iter = read_dir(entry_path.join("apps")).unwrap();
                        for icon_subdir_entry in icons_subdir_iter {
                            let icon_subdir_entry = icon_subdir_entry.unwrap();
                            install_file
                                .write_all(
                                    format!(
                                        "rm {:?}\n",
                                        icon_target_dir.join(icon_subdir_entry.file_name()),
                                    )
                                    .as_bytes(),
                                )
                                .unwrap();
                        }

                        install_file
                            .write_all(format!("rmdir -p {:?}\n", icon_target_dir).as_bytes())
                            .unwrap();
                    }
                }

                #[cfg(feature = "gettext")]
                generate_uninstall_translations_script(&mut install_file, data_dir);

                install_file
                    .write_all(b"\n# Uninstall desktop file\n")
                    .unwrap();
                let desktop_target_dir = data_dir.join("applications");
                install_file
                    .write_all(
                        format!(
                            "rm {:?}\n",
                            desktop_target_dir.join(&format!("org.fengalin.{}.desktop", *APP_NAME)),
                        )
                        .as_bytes(),
                    )
                    .unwrap();
                install_file
                    .write_all(format!("rmdir -p {:?}\n", desktop_target_dir).as_bytes())
                    .unwrap();
            }
            Err(err) => panic!("Couldn't create file `target/uninstall`: {:?}", err),
        }
    }
}

#[cfg(feature = "gettext")]
fn generate_uninstall_translations_script(install_file: &mut File, data_dir: &Path) {
    if let Ok(mut linguas_file) = File::open(&po_path().join("LINGUAS")) {
        let mut linguas = String::new();
        linguas_file
            .read_to_string(&mut linguas)
            .expect("Couldn't read po/LINGUAS as string");

        install_file
            .write_all(b"\n# Uninstall translations\n")
            .unwrap();
        let locale_base_dir = data_dir.join("locale");
        for lingua in linguas.lines() {
            let lingua_dir = locale_base_dir.join(lingua).join("LC_MESSAGES");
            install_file
                .write_all(
                    format!("rm {:?}\n", lingua_dir.join(&format!("{}.mo", *APP_NAME)),).as_bytes(),
                )
                .unwrap();

            install_file
                .write_all(format!("rmdir -p {:?}\n", lingua_dir).as_bytes())
                .unwrap();
        }
    }
}

fn main() {
    generate_resources();
    #[cfg(feature = "gettext")]
    generate_translations();

    #[cfg(target_family = "unix")]
    generate_install_script();

    #[cfg(target_family = "unix")]
    generate_uninstall_script();
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use super::stream::Stream;
use super::parser::{parse_statement, Statement, ParseError};
use super::row_builder::RowBuilder;

#[derive(Debug, PartialEq)]
pub enum DatabaseError {
    TableExists,
    QueryParseError,
}

#[derive(Debug)]
pub enum QueryResult {
    ResultSet,
    Insert(u64),
}

impl From<ParseError> for DatabaseError {
    fn from(err: ParseError) -> DatabaseError {
        DatabaseError::QueryParseError
    }
}

pub struct Database {
    tables: HashMap<String, Stream>
}

impl Database {
    pub fn new() -> Database {
        Database{
            tables: HashMap::new()
        }
    }

    pub fn create_stream(&mut self, name: &str) -> Result<&mut Stream, DatabaseError> {
        let tmp = Stream::new();
        if self.tables.contains_key(name) {
            return Err(DatabaseError::TableExists);
        }
        self.tables.insert(name.to_string(), tmp);
        let stream = self.tables.get_mut(name).unwrap();
        Ok(stream)
    }

    fn get_stream_mut(&mut self, name: &str) -> Option<&mut Stream> {
        self.tables.get_mut(name)
    }

    pub fn execute(&mut self, query: &str) -> Result<QueryResult, DatabaseError> {
        let tmp = try!(parse_statement(query));

        let result = match tmp {
            Statement::Insert(stream, row_builder) =>
                self.insert(&stream, &row_builder),
            _ => Ok(QueryResult::ResultSet)
        };
        result
    }

    pub fn insert(&mut self, query: &str, row_builder: &RowBuilder) -> Result<QueryResult, DatabaseError> {
        Ok(QueryResult::Insert(0))
    }

}

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

    // returns a valid DB for use with testing with valid simple schema
    fn get_db_with_stream() -> Database {
        let mut db = Database::new();
        db.create_stream("Jon");

        db
    }

    #[test]
    fn create_table() {
        let mut db = Database::new();
        db.create_stream("Jon");
    }

    #[test]
    fn create_table_fails_when_table_exists() {
        let mut db = get_db_with_stream();
        if let Err(result) = db.create_stream("Jon") {
            assert_eq!(result, DatabaseError::TableExists);
        } else {
            panic!("Was expecting DatabaseError::TableExists, got an OK");
        }

    }
}
<commit_msg>error checking if stream exists<commit_after>use std::collections::HashMap;
use super::stream::Stream;
use super::parser::{parse_statement, Statement, ParseError};
use super::row_builder::RowBuilder;

#[derive(Debug, PartialEq)]
pub enum DatabaseError {
    TableExists,
    QueryParseError,
    UnknownError,
    StreamNotFound,
}

#[derive(Debug)]
pub enum QueryResult {
    ResultSet,
    Insert(u64),
}

impl From<ParseError> for DatabaseError {
    fn from(err: ParseError) -> DatabaseError {
        DatabaseError::QueryParseError
    }
}

pub struct Database {
    tables: HashMap<String, Stream>
}

impl Database {
    pub fn new() -> Database {
        Database{
            tables: HashMap::new()
        }
    }

    pub fn create_stream(&mut self, name: &str) -> Result<&mut Stream, DatabaseError> {
        let tmp = Stream::new();
        if self.tables.contains_key(name) {
            return Err(DatabaseError::TableExists);
        }
        self.tables.insert(name.to_string(), tmp);
        let stream = self.tables.get_mut(name).unwrap();
        Ok(stream)
    }

    fn get_stream_mut(&mut self, name: &str) -> Option<&mut Stream> {
        self.tables.get_mut(name)
    }

    pub fn execute(&mut self, query: &str) -> Result<QueryResult, DatabaseError> {
        let tmp = try!(parse_statement(query));

        let result = match tmp {
            Statement::Insert(stream, row_builder) =>
                self.insert(&stream, &row_builder),
            _ => Err(DatabaseError::UnknownError)
        };
        result
    }

    pub fn insert(&mut self, stream: &str, row_builder: &RowBuilder) -> Result<QueryResult, DatabaseError> {
        let stream = try!(self.get_stream_mut(stream).ok_or(DatabaseError::StreamNotFound));

        Ok(QueryResult::Insert(0))
    }

}

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

    // returns a valid DB for use with testing with valid simple schema
    fn get_db_with_stream() -> Database {
        let mut db = Database::new();
        db.create_stream("Jon");

        db
    }

    #[test]
    fn create_table() {
        let mut db = Database::new();
        db.create_stream("Jon");
    }

    #[test]
    fn create_table_fails_when_table_exists() {
        let mut db = get_db_with_stream();
        if let Err(result) = db.create_stream("Jon") {
            assert_eq!(result, DatabaseError::TableExists);
        } else {
            panic!("Was expecting DatabaseError::TableExists, got an OK");
        }

    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::mem;
use std::rc::Rc;

use dense;
use error::Result;
use nfa::{self, NFA};
use sparse_set::SparseSet;
use state_id::{dead_id, StateID};

type DFARepr<S> = dense::Repr<Vec<S>, S>;

/// A determinizer converts an NFA to a DFA.
///
/// This determinizer follows the typical powerset construction, where each
/// DFA state is comprised of one or more NFA states. In the worst case, there
/// is one DFA state for every possible combination of NFA states. In practice,
/// this only happens in certain conditions, typically when there are bounded
/// repetitions.
///
/// The type variable `S` refers to the chosen state identifier representation
/// used for the DFA.
///
/// The lifetime variable `'a` refers to the lifetime of the NFA being
/// converted to a DFA.
#[derive(Debug)]
pub(crate) struct Determinizer<'a, S: StateID> {
    /// The NFA we're converting into a DFA.
    nfa: &'a NFA,
    /// The DFA we're building.
    dfa: DFARepr<S>,
    /// Each DFA state being built is defined as an *ordered* set of NFA
    /// states, along with a flag indicating whether the state is a match
    /// state or not.
    ///
    /// This is never empty. The first state is always a dummy state such that
    /// a state id == 0 corresponds to a dead state.
    builder_states: Vec<Rc<State>>,
    /// A cache of DFA states that already exist and can be easily looked up
    /// via ordered sets of NFA states.
    cache: HashMap<Rc<State>, S>,
    /// Scratch space for a stack of NFA states to visit, for depth first
    /// visiting without recursion.
    stack: Vec<nfa::StateID>,
    /// Scratch space for storing an ordered sequence of NFA states, for
    /// amortizing allocation.
    scratch_nfa_states: Vec<nfa::StateID>,
    /// Whether to build a DFA that finds the longest possible match.
    longest_match: bool,
}

/// An intermediate representation for a DFA state during determinization.
#[derive(Debug, Eq, Hash, PartialEq)]
struct State {
    /// Whether this state is a match state or not.
    is_match: bool,
    /// An ordered sequence of NFA states that make up this DFA state.
    nfa_states: Vec<nfa::StateID>,
}

impl<'a, S: StateID> Determinizer<'a, S> {
    /// Create a new determinizer for converting the given NFA to a DFA.
    pub fn new(nfa: &'a NFA) -> Determinizer<'a, S> {
        let dead = Rc::new(State::dead());
        let mut cache = HashMap::default();
        cache.insert(dead.clone(), dead_id());

        Determinizer {
            nfa,
            dfa: DFARepr::empty().anchored(nfa.is_anchored()),
            builder_states: vec![dead],
            cache,
            stack: vec![],
            scratch_nfa_states: vec![],
            longest_match: false,
        }
    }

    /// Instruct the determinizer to use equivalence classes as the transition
    /// alphabet instead of all possible byte values.
    pub fn with_byte_classes(mut self) -> Determinizer<'a, S> {
        let byte_classes = self.nfa.byte_classes().clone();
        self.dfa = DFARepr::empty_with_byte_classes(byte_classes)
            .anchored(self.nfa.is_anchored());
        self
    }

    /// Instruct the determinizer to build a DFA that recognizes the longest
    /// possible match instead of the leftmost first match. This is useful when
    /// constructing reverse DFAs for finding the start of a match.
    pub fn longest_match(mut self, yes: bool) -> Determinizer<'a, S> {
        self.longest_match = yes;
        self
    }

    /// Build the DFA. If there was a problem constructing the DFA (e.g., if
    /// the chosen state identifier representation is too small), then an error
    /// is returned.
    pub fn build(mut self) -> Result<DFARepr<S>> {
        let representative_bytes: Vec<u8> =
            self.dfa.byte_classes().representatives().collect();
        let mut sparse = self.new_sparse_set();
        let mut uncompiled = vec![self.add_start(&mut sparse)?];
        while let Some(dfa_id) = uncompiled.pop() {
            for &b in &representative_bytes {
                let (next_dfa_id, is_new) =
                    self.cached_state(dfa_id, b, &mut sparse)?;
                self.dfa.add_transition(dfa_id, b, next_dfa_id);
                if is_new {
                    uncompiled.push(next_dfa_id);
                }
            }
        }

        // At this point, we shuffle the matching states in the final DFA to
        // the beginning. This permits a DFA's match loop to detect a match
        // condition by merely inspecting the current state's identifier, and
        // avoids the need for any additional auxiliary storage.
        let is_match: Vec<bool> =
            self.builder_states.iter().map(|s| s.is_match).collect();
        self.dfa.shuffle_match_states(&is_match);
        Ok(self.dfa)
    }

    /// Return the identifier for the next DFA state given an existing DFA
    /// state and an input byte. If the next DFA state already exists, then
    /// return its identifier from the cache. Otherwise, build the state, cache
    /// it and return its identifier.
    ///
    /// The given sparse set is used for scratch space. It must have a capacity
    /// equivalent to the total number of NFA states, but its contents are
    /// otherwise unspecified.
    ///
    /// This routine returns a boolean indicating whether a new state was
    /// built. If a new state is built, then the caller needs to add it to its
    /// frontier of uncompiled DFA states to compute transitions for.
    fn cached_state(
        &mut self,
        dfa_id: S,
        b: u8,
        sparse: &mut SparseSet,
    ) -> Result<(S, bool)> {
        sparse.clear();
        // Compute the set of all reachable NFA states, including epsilons.
        self.next(dfa_id, b, sparse);
        // Build a candidate state and check if it has already been built.
        let state = self.new_state(sparse);
        if let Some(&cached_id) = self.cache.get(&state) {
            // Since we have a cached state, put the constructed state's
            // memory back into our scratch space, so that it can be reused.
            mem::replace(&mut self.scratch_nfa_states, state.nfa_states);
            return Ok((cached_id, false));
        }
        // Nothing was in the cache, so add this state to the cache.
        self.add_state(state).map(|s| (s, true))
    }

    /// Compute the set of all eachable NFA states, including the full epsilon
    /// closure, from a DFA state for a single byte of input.
    fn next(&mut self, dfa_id: S, b: u8, next_nfa_states: &mut SparseSet) {
        next_nfa_states.clear();
        for i in 0..self.builder_states[dfa_id.to_usize()].nfa_states.len() {
            let nfa_id = self.builder_states[dfa_id.to_usize()].nfa_states[i];
            match *self.nfa.state(nfa_id) {
                nfa::State::Union { .. }
                | nfa::State::Fail
                | nfa::State::Match => {}
                nfa::State::Range { range: ref r } => {
                    if r.start <= b && b <= r.end {
                        self.epsilon_closure(r.next, next_nfa_states);
                    }
                }
                nfa::State::Sparse { ref ranges } => {
                    for r in ranges.iter() {
                        if r.start > b {
                            break;
                        } else if r.start <= b && b <= r.end {
                            self.epsilon_closure(r.next, next_nfa_states);
                            break;
                        }
                    }
                }
            }
        }
    }

    /// Compute the epsilon closure for the given NFA state.
    fn epsilon_closure(&mut self, start: nfa::StateID, set: &mut SparseSet) {
        if !self.nfa.state(start).is_epsilon() {
            set.insert(start);
            return;
        }

        self.stack.push(start);
        while let Some(mut id) = self.stack.pop() {
            loop {
                if set.contains(id) {
                    break;
                }
                set.insert(id);
                match *self.nfa.state(id) {
                    nfa::State::Range { .. }
                    | nfa::State::Sparse { .. }
                    | nfa::State::Fail
                    | nfa::State::Match => break,
                    nfa::State::Union { ref alternates } => {
                        id = match alternates.get(0) {
                            None => break,
                            Some(&id) => id,
                        };
                        self.stack.extend(alternates[1..].iter().rev());
                    }
                }
            }
        }
    }

    /// Compute the initial DFA state and return its identifier.
    ///
    /// The sparse set given is used for scratch space, and must have capacity
    /// equal to the total number of NFA states. Its contents are unspecified.
    fn add_start(&mut self, sparse: &mut SparseSet) -> Result<S> {
        sparse.clear();
        self.epsilon_closure(self.nfa.start(), sparse);
        let state = self.new_state(&sparse);
        let id = self.add_state(state)?;
        self.dfa.set_start_state(id);
        Ok(id)
    }

    /// Add the given state to the DFA and make it available in the cache.
    ///
    /// The state initially has no transitions. That is, it transitions to the
    /// dead state for all possible inputs.
    fn add_state(&mut self, state: State) -> Result<S> {
        let id = self.dfa.add_empty_state()?;
        let rstate = Rc::new(state);
        self.builder_states.push(rstate.clone());
        self.cache.insert(rstate, id);
        Ok(id)
    }

    /// Convert the given set of ordered NFA states to a DFA state.
    fn new_state(&mut self, set: &SparseSet) -> State {
        let mut state = State {
            is_match: false,
            nfa_states: mem::replace(&mut self.scratch_nfa_states, vec![]),
        };
        state.nfa_states.clear();

        for &id in set {
            match *self.nfa.state(id) {
                nfa::State::Range { .. } => {
                    state.nfa_states.push(id);
                }
                nfa::State::Sparse { .. } => {
                    state.nfa_states.push(id);
                }
                nfa::State::Fail => {
                    break;
                }
                nfa::State::Match => {
                    state.is_match = true;
                    if !self.longest_match {
                        break;
                    }
                }
                nfa::State::Union { .. } => {}
            }
        }
        state
    }

    /// Create a new sparse set with enough capacity to hold all NFA states.
    fn new_sparse_set(&self) -> SparseSet {
        SparseSet::new(self.nfa.len())
    }
}

impl State {
    /// Create a new empty dead state.
    fn dead() -> State {
        State { nfa_states: vec![], is_match: false }
    }
}
<commit_msg>Silence unused_must_use warning<commit_after>use std::collections::HashMap;
use std::mem;
use std::rc::Rc;

use dense;
use error::Result;
use nfa::{self, NFA};
use sparse_set::SparseSet;
use state_id::{dead_id, StateID};

type DFARepr<S> = dense::Repr<Vec<S>, S>;

/// A determinizer converts an NFA to a DFA.
///
/// This determinizer follows the typical powerset construction, where each
/// DFA state is comprised of one or more NFA states. In the worst case, there
/// is one DFA state for every possible combination of NFA states. In practice,
/// this only happens in certain conditions, typically when there are bounded
/// repetitions.
///
/// The type variable `S` refers to the chosen state identifier representation
/// used for the DFA.
///
/// The lifetime variable `'a` refers to the lifetime of the NFA being
/// converted to a DFA.
#[derive(Debug)]
pub(crate) struct Determinizer<'a, S: StateID> {
    /// The NFA we're converting into a DFA.
    nfa: &'a NFA,
    /// The DFA we're building.
    dfa: DFARepr<S>,
    /// Each DFA state being built is defined as an *ordered* set of NFA
    /// states, along with a flag indicating whether the state is a match
    /// state or not.
    ///
    /// This is never empty. The first state is always a dummy state such that
    /// a state id == 0 corresponds to a dead state.
    builder_states: Vec<Rc<State>>,
    /// A cache of DFA states that already exist and can be easily looked up
    /// via ordered sets of NFA states.
    cache: HashMap<Rc<State>, S>,
    /// Scratch space for a stack of NFA states to visit, for depth first
    /// visiting without recursion.
    stack: Vec<nfa::StateID>,
    /// Scratch space for storing an ordered sequence of NFA states, for
    /// amortizing allocation.
    scratch_nfa_states: Vec<nfa::StateID>,
    /// Whether to build a DFA that finds the longest possible match.
    longest_match: bool,
}

/// An intermediate representation for a DFA state during determinization.
#[derive(Debug, Eq, Hash, PartialEq)]
struct State {
    /// Whether this state is a match state or not.
    is_match: bool,
    /// An ordered sequence of NFA states that make up this DFA state.
    nfa_states: Vec<nfa::StateID>,
}

impl<'a, S: StateID> Determinizer<'a, S> {
    /// Create a new determinizer for converting the given NFA to a DFA.
    pub fn new(nfa: &'a NFA) -> Determinizer<'a, S> {
        let dead = Rc::new(State::dead());
        let mut cache = HashMap::default();
        cache.insert(dead.clone(), dead_id());

        Determinizer {
            nfa,
            dfa: DFARepr::empty().anchored(nfa.is_anchored()),
            builder_states: vec![dead],
            cache,
            stack: vec![],
            scratch_nfa_states: vec![],
            longest_match: false,
        }
    }

    /// Instruct the determinizer to use equivalence classes as the transition
    /// alphabet instead of all possible byte values.
    pub fn with_byte_classes(mut self) -> Determinizer<'a, S> {
        let byte_classes = self.nfa.byte_classes().clone();
        self.dfa = DFARepr::empty_with_byte_classes(byte_classes)
            .anchored(self.nfa.is_anchored());
        self
    }

    /// Instruct the determinizer to build a DFA that recognizes the longest
    /// possible match instead of the leftmost first match. This is useful when
    /// constructing reverse DFAs for finding the start of a match.
    pub fn longest_match(mut self, yes: bool) -> Determinizer<'a, S> {
        self.longest_match = yes;
        self
    }

    /// Build the DFA. If there was a problem constructing the DFA (e.g., if
    /// the chosen state identifier representation is too small), then an error
    /// is returned.
    pub fn build(mut self) -> Result<DFARepr<S>> {
        let representative_bytes: Vec<u8> =
            self.dfa.byte_classes().representatives().collect();
        let mut sparse = self.new_sparse_set();
        let mut uncompiled = vec![self.add_start(&mut sparse)?];
        while let Some(dfa_id) = uncompiled.pop() {
            for &b in &representative_bytes {
                let (next_dfa_id, is_new) =
                    self.cached_state(dfa_id, b, &mut sparse)?;
                self.dfa.add_transition(dfa_id, b, next_dfa_id);
                if is_new {
                    uncompiled.push(next_dfa_id);
                }
            }
        }

        // At this point, we shuffle the matching states in the final DFA to
        // the beginning. This permits a DFA's match loop to detect a match
        // condition by merely inspecting the current state's identifier, and
        // avoids the need for any additional auxiliary storage.
        let is_match: Vec<bool> =
            self.builder_states.iter().map(|s| s.is_match).collect();
        self.dfa.shuffle_match_states(&is_match);
        Ok(self.dfa)
    }

    /// Return the identifier for the next DFA state given an existing DFA
    /// state and an input byte. If the next DFA state already exists, then
    /// return its identifier from the cache. Otherwise, build the state, cache
    /// it and return its identifier.
    ///
    /// The given sparse set is used for scratch space. It must have a capacity
    /// equivalent to the total number of NFA states, but its contents are
    /// otherwise unspecified.
    ///
    /// This routine returns a boolean indicating whether a new state was
    /// built. If a new state is built, then the caller needs to add it to its
    /// frontier of uncompiled DFA states to compute transitions for.
    fn cached_state(
        &mut self,
        dfa_id: S,
        b: u8,
        sparse: &mut SparseSet,
    ) -> Result<(S, bool)> {
        sparse.clear();
        // Compute the set of all reachable NFA states, including epsilons.
        self.next(dfa_id, b, sparse);
        // Build a candidate state and check if it has already been built.
        let state = self.new_state(sparse);
        if let Some(&cached_id) = self.cache.get(&state) {
            // Since we have a cached state, put the constructed state's
            // memory back into our scratch space, so that it can be reused.
            let _ =
                mem::replace(&mut self.scratch_nfa_states, state.nfa_states);
            return Ok((cached_id, false));
        }
        // Nothing was in the cache, so add this state to the cache.
        self.add_state(state).map(|s| (s, true))
    }

    /// Compute the set of all eachable NFA states, including the full epsilon
    /// closure, from a DFA state for a single byte of input.
    fn next(&mut self, dfa_id: S, b: u8, next_nfa_states: &mut SparseSet) {
        next_nfa_states.clear();
        for i in 0..self.builder_states[dfa_id.to_usize()].nfa_states.len() {
            let nfa_id = self.builder_states[dfa_id.to_usize()].nfa_states[i];
            match *self.nfa.state(nfa_id) {
                nfa::State::Union { .. }
                | nfa::State::Fail
                | nfa::State::Match => {}
                nfa::State::Range { range: ref r } => {
                    if r.start <= b && b <= r.end {
                        self.epsilon_closure(r.next, next_nfa_states);
                    }
                }
                nfa::State::Sparse { ref ranges } => {
                    for r in ranges.iter() {
                        if r.start > b {
                            break;
                        } else if r.start <= b && b <= r.end {
                            self.epsilon_closure(r.next, next_nfa_states);
                            break;
                        }
                    }
                }
            }
        }
    }

    /// Compute the epsilon closure for the given NFA state.
    fn epsilon_closure(&mut self, start: nfa::StateID, set: &mut SparseSet) {
        if !self.nfa.state(start).is_epsilon() {
            set.insert(start);
            return;
        }

        self.stack.push(start);
        while let Some(mut id) = self.stack.pop() {
            loop {
                if set.contains(id) {
                    break;
                }
                set.insert(id);
                match *self.nfa.state(id) {
                    nfa::State::Range { .. }
                    | nfa::State::Sparse { .. }
                    | nfa::State::Fail
                    | nfa::State::Match => break,
                    nfa::State::Union { ref alternates } => {
                        id = match alternates.get(0) {
                            None => break,
                            Some(&id) => id,
                        };
                        self.stack.extend(alternates[1..].iter().rev());
                    }
                }
            }
        }
    }

    /// Compute the initial DFA state and return its identifier.
    ///
    /// The sparse set given is used for scratch space, and must have capacity
    /// equal to the total number of NFA states. Its contents are unspecified.
    fn add_start(&mut self, sparse: &mut SparseSet) -> Result<S> {
        sparse.clear();
        self.epsilon_closure(self.nfa.start(), sparse);
        let state = self.new_state(&sparse);
        let id = self.add_state(state)?;
        self.dfa.set_start_state(id);
        Ok(id)
    }

    /// Add the given state to the DFA and make it available in the cache.
    ///
    /// The state initially has no transitions. That is, it transitions to the
    /// dead state for all possible inputs.
    fn add_state(&mut self, state: State) -> Result<S> {
        let id = self.dfa.add_empty_state()?;
        let rstate = Rc::new(state);
        self.builder_states.push(rstate.clone());
        self.cache.insert(rstate, id);
        Ok(id)
    }

    /// Convert the given set of ordered NFA states to a DFA state.
    fn new_state(&mut self, set: &SparseSet) -> State {
        let mut state = State {
            is_match: false,
            nfa_states: mem::replace(&mut self.scratch_nfa_states, vec![]),
        };
        state.nfa_states.clear();

        for &id in set {
            match *self.nfa.state(id) {
                nfa::State::Range { .. } => {
                    state.nfa_states.push(id);
                }
                nfa::State::Sparse { .. } => {
                    state.nfa_states.push(id);
                }
                nfa::State::Fail => {
                    break;
                }
                nfa::State::Match => {
                    state.is_match = true;
                    if !self.longest_match {
                        break;
                    }
                }
                nfa::State::Union { .. } => {}
            }
        }
        state
    }

    /// Create a new sparse set with enough capacity to hold all NFA states.
    fn new_sparse_set(&self) -> SparseSet {
        SparseSet::new(self.nfa.len())
    }
}

impl State {
    /// Create a new empty dead state.
    fn dead() -> State {
        State { nfa_states: vec![], is_match: false }
    }
}
<|endoftext|>


<commit_before>//! This crate helps to create a Fractalide agent
//!
//! It provides the macro `agent` which takes the high level view of the agent, and creates the code for the scheduler.
//!
//! It also declare the Trait Agent and the shared methods needed by every agent and the scheduler.

extern crate capnp;

// TODO : Add method to remove agents
use ports::{MsgSender, MsgReceiver};
use scheduler::Signal;
use result::Result;
use std::any::Any;

/// Provide the generic functions of agents
///
/// These three functions are used by the scheduler
pub trait Agent {
    /// Return true if there is at least one input port
    fn is_input_ports(&self) -> bool;
    /// Connect output port
    fn connect(&mut self, port: &str, sender: Box<Any + Send>) -> Result<()>;
    /// Connect array output port
    fn connect_array(&mut self, port: &str, element: String, sender: Box<Any + Send>) -> Result<()>;
    /// Add input element
    fn add_inarr_element(&mut self, port: &str, element: String, recv: Box<Any + Send>) -> Result<()>;
    /// Run the method of the agent, his personal logic
    fn run(&mut self) -> Result<Signal>;
}

/// The agent macro.
///
/// It helps to define a agent, by defining the input and output ports, if there is an option or an acc port, ...
///
/// Example :
///
/// ```rust,ignore
/// agent! {
///    inputs(input: any),
///    outputs(output: any),
///    option(prim_text),
///    fn run(&mut self) -> Result<Signal> {
///        // Receive an IP
///        let msg = try!(self.input.input.recv());
///
///        // Received an IP from the option port (a prim_text)
///        let opt = self.recv_opt();
///
///        // Get the capn'p reader
///        let reader: prim_text::Reader = try!(opt.read_schema());
///        // Print the option
///        println!("{}", try!(reader.get_text()));
///
///        // Send the received IP outside, but don't care about the success (drop on fail)
///        let _ = self.output.output.send(msg);
///
///        Ok(End)
///    }
/// }
/// ```
#[macro_export]
macro_rules! agent {
    (
        $( input($( $input_name:ident: $input_contract:ident ),*), )*
        $( inarr($( $input_a_name:ident: $input_a_contract:ident ),*), )*
        $( output($( $output_name:ident: $output_contract:ident ),*), )*
        $( outarr($( $output_a_name:ident: $output_a_contract:ident ),*), )*
        $( state( $state_type:ty => $state_value:expr ), )*
        $( option($option:ident), )*
        $( accumulator($accumulator:ident ), )*
        fn run(&mut $arg:ident) -> Result<Signal> $fun:block
    )
        =>
    {
        use rustfbp::agent::Agent;

        use rustfbp::result;
        use rustfbp::result::Result;
        use rustfbp::scheduler::{CompMsg, Signal};
        use std::error::Error;

        use std::sync::mpsc::{Sender};
        use std::sync::mpsc::channel;

        use rustfbp::ports::{MsgSender, MsgReceiver, OutputSend};

        #[allow(unused_imports)]
        use std::collections::HashMap;

        use capnp::serialize;
        use capnp::message;

        use std::io::{Read, Write};
        use std::any::Any;

        use rustfbp::scheduler::Signal::*;

        mod edge_capnp {
                include!("edge_capnp.rs");
        }
        use edge_capnp::*;

        mod edges {
                include!("edges.rs");
        }
        use edges::*;

        impl ThisAgent {
            $(
            fn dummy() {
                if stringify!($option) != "" {}
            }
                /*
            pub fn recv_option(&mut self) -> Msg {
                self.try_recv_option();
                if self.option_msg.is_none() {
                    self.option_msg = self.input.option.recv().ok();
                }
                match self.option_msg {
                    Some(ref msg) => msg.clone(),
                    None => unreachable!(),
                }
            }

            pub fn try_recv_option(&mut self) -> Option<Msg> {
                loop {
                    match self.input.option.try_recv() {
                        Err(_) => { break; },
                        Ok(msg) => { self.option_msg = Some(msg); }
                    };
                }
                self.option_msg.as_ref().map(|msg|{ msg.clone() })
            }
                */
            )*

        }

        impl Agent for ThisAgent {

            fn is_input_ports(&self) -> bool {
                $($(
                    if true || stringify!($input_name) == "" { return true; }
                )*)*
                $($(
                    if true || stringify!($input_a_name) == "" { return true; }
                )*)*
                false
            }

            fn connect(&mut self, port: &str, sender: Box<Any + Send>) -> Result<()> {
                match port {
                    $($(
                        stringify!($output_name) => {
                            let s = sender.downcast::<MsgSender<$output_contract>>().expect("cannot downast");
                            self.output.$output_name = Some(*s);
                        }
                    )*)*
                        _ => {
                            return Err(result::Error::PortDontExist(port.into()));
                        }
                }
                Ok(())
            }

            fn connect_array(&mut self, port: &str, element: String, sender: Box<Any + Send>) -> Result<()> {
                match port {
                    $($(
                        stringify!($output_a_name) => {
                            let s = sender.downcast::<MsgSender<$output_a_contract>>().expect("cannot downcast");
                            self.outarr.$output_a_name.insert(element, *s);
                        }
                    )*)*
                        _ => {
                            return Err(result::Error::PortDontExist(port.into()));
                        }
                }
                Ok(())
            }

            fn add_inarr_element(&mut self, port: &str, element: String, recv: Box<Any + Send>) -> Result<()> {
                match port {
                    $($(
                        stringify!($input_a_name) => {
                            let r = recv.downcast::<MsgReceiver<$input_a_contract>>().expect("cannot downcast");
                            self.inarr.$input_a_name.insert(element, *r);
                            Ok(())
                        }
                    )*)*
                        _ => {
                            Err(result::Error::PortDontExist(port.into()))
                        }
                }
            }

            fn run(&mut $arg) -> Result<Signal> $fun

        }

        pub struct Input {
            $($(
                $input_name: MsgReceiver<$input_contract>,
            )*)*
        }

        pub struct Inarr {
            $($(
                $input_a_name: HashMap<String, MsgReceiver<$input_a_contract>>,
            )*)*
        }

        pub struct Outarr {
            $($(
                $output_a_name: HashMap<String, MsgSender<$output_a_contract>>,
            )*)*
        }

        pub struct Output {
            $($(
                $output_name: Option<MsgSender<$output_contract>>,
            )*)*
        }

        /* Global agent */

        #[allow(dead_code)]
        #[allow(non_camel_case_types)]
        pub struct ThisAgent {
            id: usize,
            pub input: Input,
            pub output: Output,
            pub inarr: Inarr,
            pub outarr: Outarr,
            //pub option_msg: Option<Msg>,
            sched: Sender<CompMsg>,
            $(
            pub state: $state_type ,
            )*
        }

        #[allow(dead_code)]
        pub fn new(id: usize, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, Box<Any + Send>>)> {

            let mut senders: HashMap<String, Box<Any + Send>> = HashMap::new();
            // let option = MsgReceiver::new(id, sched.clone(), false);
            // senders.insert("option".to_string(), option.1);
            // let accumulator = MsgReceiver::new(id, sched.clone(), false);
            // senders.insert("accumulator".to_string(), accumulator.1.clone());
            // $($(
            //     let $input_name = MsgReceiver::new(id, sched.clone(), true);
            //     // senders.insert(stringify!($input_name).to_string(), $input_name.1);
            // )*)*

            $($(
                let $input_name = MsgReceiver::<$input_contract>::new(id, sched.clone(), true);
                senders.insert(stringify!($input_name).to_string(), Box::new($input_name.1));
            )*)*

            let input = Input {
                $($(
                    $input_name: $input_name.0,
                )*)*
            };
            let output = Output {
                $($(
                    $output_name: None,
                )*)*
            };
            let inarr = Inarr {
                $($(
                    $input_a_name: HashMap::new(),
                )*)*
            };
            let outarr = Outarr {
                $($(
                    $output_a_name: HashMap::new(),
                )*)*
            };

            let agent= ThisAgent {
                id: id,
                input: input,
                output: output,
                inarr: inarr,
                outarr: outarr,
                //option_msg: None,
                sched: sched,
                $(
                    state: $state_value,
                )*
            };

            Ok((Box::new(agent) as Box<Agent + Send>, senders))
        }

        #[no_mangle]
        pub extern fn create_agent(id: usize, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, Box<Any + Send>>)> {
            new(id, sched)
        }

        #[no_mangle]
        pub extern fn clone_input(port: &str, sender: &Box<Any + Send>) -> Result<Box<Any + Send>> {
            match port {
                $($(
                    stringify!($input_name) => {
                        let s = sender.downcast_ref::<MsgSender<$input_contract>>().unwrap();
                        Ok(Box::new(s.clone()))
                    },
                )*)*
                    _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn clone_input_array(port: &str, sender: &Box<Any + Send>) -> Result<Box<Any + Send>> {
            match port {
                $($(
                    stringify!($input_a_name) => {
                        let s = sender.downcast_ref::<MsgSender<$input_a_contract>>().unwrap();
                        Ok(Box::new(s.clone()))
                    },
                )*)*
                    _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn create_input_array(port: &str, id: usize, sched: Sender<CompMsg>, must_sched: bool ) -> Result<(Box<Any + Send>, Box<Any + Send>)> {
            match port {
                $($(
                    stringify!($input_a_name) => {
                        let (r, s): (MsgReceiver::<$input_a_contract>, MsgSender::<$input_a_contract>) = MsgReceiver::new(id, sched, must_sched);
                        Ok((Box::new(r), Box::new(s)))
                    },
                )*)*
                    _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_input(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($input_name)=> Ok(stringify!($input_contract).into()),
                )*)*
                /*
                $(
                    "option" => Ok(stringify!($option).into()),
                )*
                $(
                    "accumulator" => Ok(stringify!($accumulator).into()),
                )*
                */
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_input_array(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($input_a_name) => Ok(stringify!($input_a_contract).into()),
                )*)*
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($output_name)=> Ok(stringify!($output_contract).into()),
                )*)*
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output_array(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($output_a_name) => Ok(stringify!($output_a_contract).into()),
                )*)*
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }
    }
}

#[macro_export]
macro_rules! send_action {
    ($agent: ident, $port:ident, $msg:ident) => {{
        if let Some(sender) = $agent.outarr.$port.get(&$msg.action) {
            sender.send($msg)
        } else {
            $agent.output.$port.send($msg)
        }
    }}
}
<commit_msg>Acc, Option and state are coming back<commit_after>//! This crate helps to create a Fractalide agent
//!
//! It provides the macro `agent` which takes the high level view of the agent, and creates the code for the scheduler.
//!
//! It also declare the Trait Agent and the shared methods needed by every agent and the scheduler.

extern crate capnp;

// TODO : Add method to remove agents
use ports::{MsgSender, MsgReceiver};
use scheduler::Signal;
use result::Result;
use std::any::Any;

/// Provide the generic functions of agents
///
/// These three functions are used by the scheduler
pub trait Agent {
    /// Return true if there is at least one input port
    fn is_input_ports(&self) -> bool;
    /// Connect output port
    fn connect(&mut self, port: &str, sender: Box<Any + Send>) -> Result<()>;
    /// Connect array output port
    fn connect_array(&mut self, port: &str, element: String, sender: Box<Any + Send>) -> Result<()>;
    /// Add input element
    fn add_inarr_element(&mut self, port: &str, element: String, recv: Box<Any + Send>) -> Result<()>;
    /// Run the method of the agent, his personal logic
    fn run(&mut self) -> Result<Signal>;
}

/// The agent macro.
///
/// It helps to define a agent, by defining the input and output ports, if there is an option or an acc port, ...
///
/// Example :
///
/// ```rust,ignore
/// agent! {
///    inputs(input: any),
///    outputs(output: any),
///    option(prim_text),
///    fn run(&mut self) -> Result<Signal> {
///        // Receive an IP
///        let msg = try!(self.input.input.recv());
///
///        // Received an IP from the option port (a prim_text)
///        let opt = self.recv_opt();
///
///        // Get the capn'p reader
///        let reader: prim_text::Reader = try!(opt.read_schema());
///        // Print the option
///        println!("{}", try!(reader.get_text()));
///
///        // Send the received IP outside, but don't care about the success (drop on fail)
///        let _ = self.output.output.send(msg);
///
///        Ok(End)
///    }
/// }
/// ```
#[macro_export]
macro_rules! agent {
    (
        $( input($( $input_name:ident: $input_contract:ident ),*), )*
        $( inarr($( $input_a_name:ident: $input_a_contract:ident ),*), )*
        $( output($( $output_name:ident: $output_contract:ident ),*), )*
        $( outarr($( $output_a_name:ident: $output_a_contract:ident ),*), )*
        $( state( $state_type:ty => $state_value:expr ), )*
        $( option($option:ident), )*
        $( accumulator($accumulator:ident ), )*
        fn run(&mut $arg:ident) -> Result<Signal> $fun:block
    )
        =>
    {
        use rustfbp::agent::Agent;

        use rustfbp::result;
        use rustfbp::result::Result;
        use rustfbp::scheduler::{CompMsg, Signal};
        use std::error::Error;

        use std::sync::mpsc::{Sender};
        use std::sync::mpsc::channel;

        use rustfbp::ports::{MsgSender, MsgReceiver, OutputSend};

        #[allow(unused_imports)]
        use std::collections::HashMap;

        use capnp::serialize;
        use capnp::message;

        use std::io::{Read, Write};
        use std::any::Any;

        use rustfbp::scheduler::Signal::*;

        mod edge_capnp {
                include!("edge_capnp.rs");
        }
        use edge_capnp::*;

        mod edges {
                include!("edges.rs");
        }
        use edges::*;

        impl ThisAgent {
            $(

            pub fn recv_option(&mut self) -> $option {
                self.try_recv_option();
                if self.option_msg.is_none() {
                    self.option_msg = self.input.option.recv().ok();
                }
                match self.option_msg {
                    Some(ref msg) => msg.clone(),
                    None => unreachable!(),
                }
            }

            pub fn try_recv_option(&mut self) -> Option<$option> {
                loop {
                    match self.input.option.try_recv() {
                        Err(_) => { break; },
                        Ok(msg) => { self.option_msg = Some(msg); }
                    };
                }
                self.option_msg.as_ref().map(|msg|{ msg.clone() })
            }
            )*

        }

        impl Agent for ThisAgent {

            fn is_input_ports(&self) -> bool {
                $($(
                    if true || stringify!($input_name) == "" { return true; }
                )*)*
                $($(
                    if true || stringify!($input_a_name) == "" { return true; }
                )*)*
                false
            }

            fn connect(&mut self, port: &str, sender: Box<Any + Send>) -> Result<()> {
                match port {
                    $($(
                        stringify!($output_name) => {
                            let s = sender.downcast::<MsgSender<$output_contract>>().expect("cannot downast");
                            self.output.$output_name = Some(*s);
                        }
                    )*)*
                        _ => {
                            return Err(result::Error::PortDontExist(port.into()));
                        }
                }
                Ok(())
            }

            fn connect_array(&mut self, port: &str, element: String, sender: Box<Any + Send>) -> Result<()> {
                match port {
                    $($(
                        stringify!($output_a_name) => {
                            let s = sender.downcast::<MsgSender<$output_a_contract>>().expect("cannot downcast");
                            self.outarr.$output_a_name.insert(element, *s);
                        }
                    )*)*
                        _ => {
                            return Err(result::Error::PortDontExist(port.into()));
                        }
                }
                Ok(())
            }

            fn add_inarr_element(&mut self, port: &str, element: String, recv: Box<Any + Send>) -> Result<()> {
                match port {
                    $($(
                        stringify!($input_a_name) => {
                            let r = recv.downcast::<MsgReceiver<$input_a_contract>>().expect("cannot downcast");
                            self.inarr.$input_a_name.insert(element, *r);
                            Ok(())
                        }
                    )*)*
                        _ => {
                            Err(result::Error::PortDontExist(port.into()))
                        }
                }
            }

            fn run(&mut $arg) -> Result<Signal> $fun

        }

        pub struct Input {
            $($(
                $input_name: MsgReceiver<$input_contract>,
            )*)*
            $(
                option: MsgReceiver<$option>,
            )*
            $(
                accumulator: MsgReceiver<$accumulator>,
            )*
        }

        pub struct Inarr {
            $($(
                $input_a_name: HashMap<String, MsgReceiver<$input_a_contract>>,
            )*)*
        }

        pub struct Outarr {
            $($(
                $output_a_name: HashMap<String, MsgSender<$output_a_contract>>,
            )*)*
        }

        pub struct Output {
            $($(
                $output_name: Option<MsgSender<$output_contract>>,
            )*)*
            $(
                accumulator: Option<MsgSender<$accumulator>>,
            )*
        }

        /* Global agent */

        #[allow(dead_code)]
        #[allow(non_camel_case_types)]
        pub struct ThisAgent {
            id: usize,
            pub input: Input,
            pub output: Output,
            pub inarr: Inarr,
            pub outarr: Outarr,
            $(
                pub option_msg: Option<$option>,
            )*
            sched: Sender<CompMsg>,
            $(
            pub state: $state_type ,
            )*
        }

        #[allow(dead_code)]
        pub fn new(id: usize, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, Box<Any + Send>>)> {

            let mut senders: HashMap<String, Box<Any + Send>> = HashMap::new();
            $(
                let option = MsgReceiver::<$option>::new(id, sched.clone(), false);
                senders.insert("option".to_string(), Box::new(option.1));
            )*

            $(
                let accumulator = MsgReceiver::<$accumulator>::new(id, sched.clone(), false);
                senders.insert("accumulator".to_string(), Box::new(accumulator.1.clone()));
            )*

            $($(
                let $input_name = MsgReceiver::<$input_contract>::new(id, sched.clone(), true);
                senders.insert(stringify!($input_name).to_string(), Box::new($input_name.1));
            )*)*

            let input = Input {
                $($(
                    $input_name: $input_name.0,
                )*)*
                $(
                    option: option.0 as MsgReceiver::<$option>,
                )*
                $(
                    accumulator: accumulator.0 as MsgReceiver::<$accumulator>,
                )*
            };
            let output = Output {
                $($(
                    $output_name: None,
                )*)*
                $(
                    accumulator: Some(accumulator.1) as Option<MsgSender::<$accumulator>>,
                )*
            };
            let inarr = Inarr {
                $($(
                    $input_a_name: HashMap::new(),
                )*)*
            };
            let outarr = Outarr {
                $($(
                    $output_a_name: HashMap::new(),
                )*)*
            };

            let agent= ThisAgent {
                id: id,
                input: input,
                output: output,
                inarr: inarr,
                outarr: outarr,
                $(
                    option_msg: None as Option<$option>,
                )*
                sched: sched,
                $(
                    state: $state_value,
                )*
            };

            Ok((Box::new(agent) as Box<Agent + Send>, senders))
        }

        #[no_mangle]
        pub extern fn create_agent(id: usize, sched: Sender<CompMsg>) -> Result<(Box<Agent + Send>, HashMap<String, Box<Any + Send>>)> {
            new(id, sched)
        }

        #[no_mangle]
        pub extern fn clone_input(port: &str, sender: &Box<Any + Send>) -> Result<Box<Any + Send>> {
            match port {
                $($(
                    stringify!($input_name) => {
                        let s = sender.downcast_ref::<MsgSender<$input_contract>>().unwrap();
                        Ok(Box::new(s.clone()))
                    },
                )*)*
                    $(
                        "option" => {
                            let s = sender.downcast_ref::<MsgSender<$option>>().unwrap();
                            Ok(Box::new(s.clone()))
                        }
                    )*
                    $(
                        "accumulator" => {
                            let s = sender.downcast_ref::<MsgSender<$accumulator>>().unwrap();
                            Ok(Box::new(s.clone()))
                        }
                    )*
                    _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn clone_input_array(port: &str, sender: &Box<Any + Send>) -> Result<Box<Any + Send>> {
            match port {
                $($(
                    stringify!($input_a_name) => {
                        let s = sender.downcast_ref::<MsgSender<$input_a_contract>>().unwrap();
                        Ok(Box::new(s.clone()))
                    },
                )*)*
                    _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn create_input_array(port: &str, id: usize, sched: Sender<CompMsg>, must_sched: bool ) -> Result<(Box<Any + Send>, Box<Any + Send>)> {
            match port {
                $($(
                    stringify!($input_a_name) => {
                        let (r, s): (MsgReceiver::<$input_a_contract>, MsgSender::<$input_a_contract>) = MsgReceiver::new(id, sched, must_sched);
                        Ok((Box::new(r), Box::new(s)))
                    },
                )*)*
                    _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_input(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($input_name)=> Ok(stringify!($input_contract).into()),
                )*)*
                $(
                    "option" => Ok(stringify!($option).into()),
                )*
                $(
                    "accumulator" => Ok(stringify!($accumulator).into()),
                )*
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_input_array(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($input_a_name) => Ok(stringify!($input_a_contract).into()),
                )*)*
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($output_name)=> Ok(stringify!($output_contract).into()),
                )*)*
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }

        #[no_mangle]
        pub extern fn get_schema_output_array(port: &str) -> Result<String> {
            match port {
                $($(
                    stringify!($output_a_name) => Ok(stringify!($output_a_contract).into()),
                )*)*
                _ => { Err(result::Error::PortDontExist(port.into())) }
            }
        }
    }
}

#[macro_export]
macro_rules! send_action {
    ($agent: ident, $port:ident, $msg:ident) => {{
        if let Some(sender) = $agent.outarr.$port.get(&$msg.action) {
            sender.send($msg)
        } else {
            $agent.output.$port.send($msg)
        }
    }}
}
<|endoftext|>


<commit_before>use std::mem;
use std::f64;

use itertools::Itertools;
use ndarray::prelude::*;
use bit_vec::BitVec;
use ordered_float::NotNaN;
use ndarray_rand::RandomExt;
use rand::distributions::Range;

use bio::stats::{Prob, LogProb};

use io::codebook::Codebook;

pub type Readout = BitVec;

pub fn estimate<I: Iterator<Item=(String, Readout)>>(
    codebook: &Codebook,
    readouts: I
) -> (Vec<Prob>, Vec<Prob>) {
    let max_diff = 0.001;
    let shape = (codebook.len(), codebook.N as usize);

    let mut observed_01_errors = Array2::from_elem(shape, 0.0);
    let mut observed_10_errors = Array2::from_elem(shape, 0.0);
    let mut feat_count = Array2::from_elem((codebook.len(), 1), 0.0);
    for (feat, readout) in readouts {
        let feature_id = codebook.get_id(&feat);
        let encoding = codebook.record(feature_id).codeword();

        // count features
        feat_count[(feature_id, 0)] += 1.0;

        for k in 0..codebook.N as usize {
            if encoding[k] == false {
                // count 01 errors
                if readout[k] == true {
                    observed_01_errors[(feature_id, k)] += 1.0;
                }
            } else {
                // count 10 errors
                if readout[k] == false {
                    observed_10_errors[(feature_id, k)] += 1.0;
                }
            }
        }
    }

    // calculate observed (joint) error rates
    let observed_p0 = observed_01_errors / &feat_count;
    let observed_p1 = observed_10_errors / &feat_count;

    // run EM algorithm to estimate error rates
    let mut prob_no_error = Array2::from_elem(
        (codebook.len(), codebook.N as usize), LogProb::ln_one()
    );

    let mut p0 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));
    let mut p1 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));

    let mut p0_last = Array1::from_elem(codebook.N as usize, 0.0);
    let mut p1_last = Array1::from_elem(codebook.N as usize, 0.0);
    for i in 0..1000 {
        // E-step: estimate probability to have no errors in rest of readout
        for feature_id in 0..codebook.len() {
            let encoding = codebook.record(feature_id).codeword();
            for k in 0..codebook.N as usize {
                let p = LogProb((0..codebook.N as usize).filter_map(|l| {
                    if l != k {
                        let p_err = if encoding[l] == false {
                            p0[k]
                        } else {
                            p1[k]
                        };
                        Some(*LogProb((p_err as f64).ln()).ln_one_minus_exp())
                    } else { None }
                }).sum());
                prob_no_error[(feature_id, k)] = p;
            }
        }

        // M-step: find p0 and p1 that maximizes the likelihood
        mem::swap(&mut p0, &mut p0_last);
        mem::swap(&mut p1, &mut p1_last);
        for k in 0..codebook.N as usize {
            // solve least squares: take mean of all estimates
            let p = |observed_p: &Array2<f64>, bit: bool| {
                let mut count = 0.0;
                let p = (0..codebook.len()).filter_map(|feature_id| {
                    let encoding = codebook.record(feature_id).codeword();
                    if feat_count[(feature_id, 0)] < 50.0 || encoding[k] != bit {
                        None
                    } else {
                        count += 1.0;
                        let p = observed_p[(feature_id, k)] * prob_no_error[(feature_id, k)].exp();
                        Some(p)
                    }
                }).sum::<f64>() / count;
                if count == 0.0 {
                    panic!("not enough counts to estimate error rate");
                }
                p
            };
            //debug!("----------- p0 ------------");
            p0[k] = p(&observed_p0, false);
            //debug!("DEBUG: p0={}", p0[k]);
            //debug!("----------- p1 ------------");
            p1[k] = p(&observed_p1, true);
            //debug!("DEBUG: p1={}", p1[k]);
        }

        // check convergence
        let abs = |d: &f64| NotNaN::new(d.abs()).unwrap();
        //debug!("{:?}", p0);
        let convergence = p0.all_close(&p0_last, max_diff) && p1.all_close(&p1_last, max_diff);

        if convergence && i >= 5 {
            debug!("convergence reached");
            debug!("{:?}", p1);
            break;
        }

        debug!("EM-algorithm step {} of {}.", i + 1, 100);
    }

    (
        p0.iter().map(|p| Prob(*p)).collect_vec(),
        p1.iter().map(|p| Prob(*p)).collect_vec()
    )
}
<commit_msg>Better error message.<commit_after>use std::mem;
use std::f64;

use itertools::Itertools;
use ndarray::prelude::*;
use bit_vec::BitVec;
use ordered_float::NotNaN;
use ndarray_rand::RandomExt;
use rand::distributions::Range;

use bio::stats::{Prob, LogProb};

use io::codebook::Codebook;

pub type Readout = BitVec;

pub fn estimate<I: Iterator<Item=(String, Readout)>>(
    codebook: &Codebook,
    readouts: I
) -> (Vec<Prob>, Vec<Prob>) {
    let max_diff = 0.001;
    let shape = (codebook.len(), codebook.N as usize);

    let mut observed_01_errors = Array2::from_elem(shape, 0.0);
    let mut observed_10_errors = Array2::from_elem(shape, 0.0);
    let mut feat_count = Array2::from_elem((codebook.len(), 1), 0.0);
    for (feat, readout) in readouts {
        let feature_id = codebook.get_id(&feat);
        let encoding = codebook.record(feature_id).codeword();

        // count features
        feat_count[(feature_id, 0)] += 1.0;

        for k in 0..codebook.N as usize {
            if encoding[k] == false {
                // count 01 errors
                if readout[k] == true {
                    observed_01_errors[(feature_id, k)] += 1.0;
                }
            } else {
                // count 10 errors
                if readout[k] == false {
                    observed_10_errors[(feature_id, k)] += 1.0;
                }
            }
        }
    }

    // calculate observed (joint) error rates
    let observed_p0 = observed_01_errors / &feat_count;
    let observed_p1 = observed_10_errors / &feat_count;

    // run EM algorithm to estimate error rates
    let mut prob_no_error = Array2::from_elem(
        (codebook.len(), codebook.N as usize), LogProb::ln_one()
    );

    let mut p0 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));
    let mut p1 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));

    let mut p0_last = Array1::from_elem(codebook.N as usize, 0.0);
    let mut p1_last = Array1::from_elem(codebook.N as usize, 0.0);
    for i in 0..1000 {
        // E-step: estimate probability to have no errors in rest of readout
        for feature_id in 0..codebook.len() {
            let encoding = codebook.record(feature_id).codeword();
            for k in 0..codebook.N as usize {
                let p = LogProb((0..codebook.N as usize).filter_map(|l| {
                    if l != k {
                        let p_err = if encoding[l] == false {
                            p0[k]
                        } else {
                            p1[k]
                        };
                        Some(*LogProb((p_err as f64).ln()).ln_one_minus_exp())
                    } else { None }
                }).sum());
                prob_no_error[(feature_id, k)] = p;
            }
        }

        // M-step: find p0 and p1 that maximizes the likelihood
        mem::swap(&mut p0, &mut p0_last);
        mem::swap(&mut p1, &mut p1_last);
        for k in 0..codebook.N as usize {
            // solve least squares: take mean of all estimates
            let p = |observed_p: &Array2<f64>, bit: bool| {
                let mut count = 0.0;
                let p = (0..codebook.len()).filter_map(|feature_id| {
                    let encoding = codebook.record(feature_id).codeword();
                    if feat_count[(feature_id, 0)] < 50.0 || encoding[k] != bit {
                        None
                    } else {
                        count += 1.0;
                        let p = observed_p[(feature_id, k)] * prob_no_error[(feature_id, k)].exp();
                        Some(p)
                    }
                }).sum::<f64>() / count;
                if count == 0.0 {
                    panic!("not enough counts to estimate error rate (at least 1 feature with 50 counts over all cells required)");
                }
                p
            };
            //debug!("----------- p0 ------------");
            p0[k] = p(&observed_p0, false);
            //debug!("DEBUG: p0={}", p0[k]);
            //debug!("----------- p1 ------------");
            p1[k] = p(&observed_p1, true);
            //debug!("DEBUG: p1={}", p1[k]);
        }

        // check convergence
        let abs = |d: &f64| NotNaN::new(d.abs()).unwrap();
        //debug!("{:?}", p0);
        let convergence = p0.all_close(&p0_last, max_diff) && p1.all_close(&p1_last, max_diff);

        if convergence && i >= 5 {
            debug!("convergence reached");
            debug!("{:?}", p1);
            break;
        }

        debug!("EM-algorithm step {} of {}.", i + 1, 100);
    }

    (
        p0.iter().map(|p| Prob(*p)).collect_vec(),
        p1.iter().map(|p| Prob(*p)).collect_vec()
    )
}
<|endoftext|>


<commit_before>macro_rules! u8_enum {
    {
        $name: ident {
            $($field_name: ident = $value: expr),+,
        }
    } => {
        #[derive(PartialEq, Eq, Clone, Copy, Debug)]
        pub enum $name {
            $(
                $field_name = $value,
            )+
        }

        impl Into<u8> for $name {
            fn into(self) -> u8 {
                self as u8
            }
        }

        impl From<u8> for $name {
            fn from(v: u8) -> $name {
                match v {
                    $(
                        $value => $name::$field_name
                    ),+,
                    _ => panic!("Enum value not recognized: ${:02X}", v),
                }
            }
        }
    }
}

macro_rules! memory_mapper {
    {
        name: $name: ident,
        fields: [
            $($hex_f:expr, $field_name_f: ident, $default_f: expr);*;
        ],
        bitfields: {
            getters: [
                $($hex:expr, $field_name: ident, $default: expr, [
                    $($bitfield_getter: ident,
                    $method_getter: ident,
                    $method_type: ident);*
                ]);*
            ],
            getter_setters: [
                $($hex_s:expr, $field_name_s: ident, $default_s: expr, [
                    $($bitfield_getter_s: ident,
                    $bitfield_setter_s: ident,
                    $method_getter_s: ident,
                    $method_setter_s: ident,
                    $method_type_s: ident);*
                ]);*
            ],
        },
    } => {
        struct $name {
            $($field_name: Bitfield),+,
            $($field_name_f: u8),+
            $(,$field_name_s: Bitfield),*
        }

        #[allow(dead_code)]
        impl $name {
            pub fn new() -> $name {
                $name {
                    $($field_name: Bitfield::new($default)),+,
                    $($field_name_f: $default_f),+,
                    $($field_name_s: Bitfield::new($default_s)),*
                }
            }

            $($(
                pub fn $method_getter(&self) -> $method_type {
                    self.$field_name.$bitfield_getter().into()
                }
            )*)*

            $($(
                pub fn $method_getter_s(&self) -> $method_type_s {
                    self.$field_name_s.$bitfield_getter_s().into()
                }

                pub fn $method_setter_s(&mut self, v: $method_type_s) {
                    self.$field_name_s.$bitfield_setter_s(v.into());
                }
            )*)*
        }

        impl Handler for $name {
            fn read(&self, address: u16) -> u8 {
                match address {
                    $($hex => self.$field_name.get()),+,
                    $($hex_f => self.$field_name_f),+
                    $(,$hex_s => self.$field_name_s.get()),*,
                    _ => panic!("Could not handle read at ${:04X}", address),
                }
            }

            fn write(&mut self, address: u16, v: u8) {
                match address {
                    $($hex_f => self.$field_name_f = v),+,
                    $($hex => {
                        self.$field_name.set(v);
                    }),+
                    $(,$hex_s => self.$field_name_s.set(v)),*,
                    _ => panic!("Could not handle write at ${:04X} v=${:02X}", address, v),
                }
            }
        }
    }
}

macro_rules! memory_handler {
    {
        parent: $parent: ident,
        mapper: $mapper: ident,
        callback: $callback: ident,
    } => {
        impl Handler for $parent {
            fn read(&self, address: u16) -> u8 {
                self.$mapper.read(address)
            }

            fn write(&mut self, address: u16, v: u8) {
                self.$mapper.write(address, v);
                self.$callback(address);
            }
        }
    }
}

#[allow(non_snake_case)]
pub mod opcodes;
#[allow(non_snake_case)]
pub mod cpu;

#[allow(non_snake_case)]
#[allow(non_camel_case_types)]
pub mod memory_controller;
pub mod video_controller;
pub mod timer_controller;
pub mod sound_controller;

pub mod cartridge;

#[allow(non_snake_case)]
pub mod handler_holder;
<commit_msg>Fix bug in handling of params for memory_mapper!<commit_after>macro_rules! u8_enum {
    {
        $name: ident {
            $($field_name: ident = $value: expr),+,
        }
    } => {
        #[derive(PartialEq, Eq, Clone, Copy, Debug)]
        pub enum $name {
            $(
                $field_name = $value,
            )+
        }

        impl Into<u8> for $name {
            fn into(self) -> u8 {
                self as u8
            }
        }

        impl From<u8> for $name {
            fn from(v: u8) -> $name {
                match v {
                    $(
                        $value => $name::$field_name
                    ),+,
                    _ => panic!("Enum value not recognized: ${:02X}", v),
                }
            }
        }
    }
}

macro_rules! memory_mapper {
    {
        name: $name: ident,
        fields: [
            $($hex_f:expr, $field_name_f: ident, $default_f: expr);*;
        ],
        bitfields: {
            getters: [
                $($hex:expr, $field_name: ident, $default: expr, [
                    $($bitfield_getter: ident,
                    $method_getter: ident,
                    $method_type: ident);*
                ]);*
            ],
            getter_setters: [
                $($hex_s:expr, $field_name_s: ident, $default_s: expr, [
                    $($bitfield_getter_s: ident,
                    $bitfield_setter_s: ident,
                    $method_getter_s: ident,
                    $method_setter_s: ident,
                    $method_type_s: ident);*
                ]);*
            ],
        },
    } => {
        struct $name {
            $($field_name: Bitfield),+,
            $($field_name_f: u8),+
            $(,$field_name_s: Bitfield)*
        }

        #[allow(dead_code)]
        impl $name {
            pub fn new() -> $name {
                $name {
                    $($field_name: Bitfield::new($default)),+,
                    $($field_name_f: $default_f),+,
                    $($field_name_s: Bitfield::new($default_s)),*
                }
            }

            $($(
                pub fn $method_getter(&self) -> $method_type {
                    self.$field_name.$bitfield_getter().into()
                }
            )*)*

            $($(
                pub fn $method_getter_s(&self) -> $method_type_s {
                    self.$field_name_s.$bitfield_getter_s().into()
                }

                pub fn $method_setter_s(&mut self, v: $method_type_s) {
                    self.$field_name_s.$bitfield_setter_s(v.into());
                }
            )*)*
        }

        impl Handler for $name {
            fn read(&self, address: u16) -> u8 {
                match address {
                    $($hex => self.$field_name.get()),+,
                    $($hex_f => self.$field_name_f),+
                    $(,$hex_s => self.$field_name_s.get())*,
                    _ => panic!("Could not handle read at ${:04X}", address),
                }
            }

            fn write(&mut self, address: u16, v: u8) {
                match address {
                    $($hex_f => self.$field_name_f = v),+,
                    $($hex => {
                        self.$field_name.set(v);
                    }),+
                    $(,$hex_s => self.$field_name_s.set(v))*,
                    _ => panic!("Could not handle write at ${:04X} v=${:02X}", address, v),
                }
            }
        }
    }
}

macro_rules! memory_handler {
    {
        parent: $parent: ident,
        mapper: $mapper: ident,
        callback: $callback: ident,
    } => {
        impl Handler for $parent {
            fn read(&self, address: u16) -> u8 {
                self.$mapper.read(address)
            }

            fn write(&mut self, address: u16, v: u8) {
                self.$mapper.write(address, v);
                self.$callback(address);
            }
        }
    }
}

#[allow(non_snake_case)]
pub mod opcodes;
#[allow(non_snake_case)]
pub mod cpu;

#[allow(non_snake_case)]
#[allow(non_camel_case_types)]
pub mod memory_controller;
pub mod video_controller;
pub mod timer_controller;
pub mod sound_controller;

pub mod cartridge;

#[allow(non_snake_case)]
pub mod handler_holder;
<|endoftext|>