<commit_before>extern crate chrono;
extern crate clap;
extern crate diesel;
extern crate dotenv;

mod database_error;
#[macro_use]
mod database;

use chrono::*;
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};
#[cfg(feature = "postgres")]
use diesel::pg::PgConnection;
#[cfg(feature = "sqlite")]
use diesel::sqlite::SqliteConnection;
use diesel::migrations::schema::*;
use diesel::types::{FromSql, VarChar};
use diesel::{migrations, Connection, Insertable};
use std::error::Error;
use std::io::stdout;
use std::path::{PathBuf, Path};
use std::{env, fs};

use self::database_error::{DatabaseError, DatabaseResult};

fn main() {
    use self::dotenv::dotenv;
    dotenv().ok();

    let database_arg = Arg::with_name("DATABASE_URL")
        .long("database-url")
        .help("Specifies the database URL to connect to. Falls back to \
                   the DATABASE_URL environment variable if unspecified.")
        .global(true)
        .takes_value(true);

    let migration_subcommand = SubCommand::with_name("migration")
        .setting(AppSettings::VersionlessSubcommands)
        .arg(Arg::with_name("MIGRATION_DIRECTORY")
            .long("migration-dir")
            .help("The location of your migration directory. By default this \
                   will look for a directory called `migrations` in the \
                   current directory and its parents.")
            .takes_value(true)
            .global(true)
        ).subcommand(
            SubCommand::with_name("run")
                .about("Runs all pending migrations")
        ).subcommand(
            SubCommand::with_name("revert")
                .about("Reverts the latest run migration")
        ).subcommand(
            SubCommand::with_name("redo")
                .about("Reverts and re-runs the latest migration. Useful \
                      for testing that a migration can in fact be reverted.")
        ).subcommand(
            SubCommand::with_name("generate")
                .about("Generate a new migration with the given name, and \
                      the current timestamp as the version")
                .arg(Arg::with_name("MIGRATION_NAME")
                     .help("The name of the migration to create")
                     .required(true)
                 )
                .arg(Arg::with_name("MIGRATION_VERSION")
                     .long("version")
                     .help("The version number to use when generating the migration. \
                            Defaults to the current timestamp, which should suffice \
                            for most use cases.")
                     .takes_value(true)
                )
        ).setting(AppSettings::SubcommandRequiredElseHelp);

    let setup_subcommand = SubCommand::with_name("setup")
        .about("Creates the migrations directory, creates the database \
                specified in your DATABASE_URL, and runs existing migrations.");

    let database_subcommand = SubCommand::with_name("database")
        .setting(AppSettings::VersionlessSubcommands)
        .subcommand(
            SubCommand::with_name("setup")
                .about("Creates the migrations directory, creates the database \
                        specified in your DATABASE_URL, and runs existing migrations.")
        ).subcommand(
            SubCommand::with_name("reset")
                .about("Resets your database by dropping the database specified \
                        in your DATABASE_URL and then running `diesel database setup`.")
        ).subcommand(
            SubCommand::with_name("drop")
                .about("Drops the database specified in your DATABASE_URL")
                .setting(AppSettings::Hidden)
        ).setting(AppSettings::SubcommandRequiredElseHelp);

    let matches = App::new("diesel")
        .version(env!("CARGO_PKG_VERSION"))
        .setting(AppSettings::VersionlessSubcommands)
        .arg(database_arg)
        .subcommand(migration_subcommand)
        .subcommand(setup_subcommand)
        .subcommand(database_subcommand)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .get_matches();

    match matches.subcommand() {
        ("migration", Some(matches)) => run_migration_command(matches),
        ("setup", Some(matches)) => run_setup_command(matches),
        ("database", Some(matches)) => run_database_command(matches),
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

fn run_migration_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("run", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, migrations::run_pending_migrations)
                .unwrap_or_else(handle_error);
        }
        ("revert", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, migrations::revert_latest_migration)
                .unwrap_or_else(handle_error);
        }
        ("redo", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, redo_latest_migration);
        }
        ("generate", Some(args)) => {
            let migration_name = args.value_of("MIGRATION_NAME").unwrap();
            let version = migration_version(args);
            let versioned_name = format!("{}_{}", version, migration_name);
            let migration_dir = migrations_dir(args).join(versioned_name);
            fs::create_dir(&migration_dir).unwrap();

            let migration_dir_relative = convert_absolute_path_to_relative(
                &migration_dir,
                &env::current_dir().unwrap()
            );

            let up_path = migration_dir.join("up.sql");
            println!("Creating {}", migration_dir_relative.join("up.sql").display());
            fs::File::create(up_path).unwrap();
            let down_path = migration_dir.join("down.sql");
            println!("Creating {}", migration_dir_relative.join("down.sql").display());
            fs::File::create(down_path).unwrap();
        }
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

use std::fmt::Display;
fn migration_version<'a>(matches: &'a ArgMatches) -> Box<Display + 'a> {
    matches.value_of("MIGRATION_VERSION").map(|s| Box::new(s) as Box<Display>)
        .unwrap_or_else(|| Box::new(Local::now().format("%Y%m%d%H%M%S")))
}

fn migrations_dir(matches: &ArgMatches) -> PathBuf {
    matches.value_of("MIGRATION_DIRECTORY")
        .map(PathBuf::from)
        .or_else(|| {
            env::var("MIGRATION_DIRECTORY").map(PathBuf::from).ok()
        }).unwrap_or_else(|| {
            migrations::find_migrations_directory()
                .unwrap_or_else(handle_error)
        })
}

fn run_setup_command(matches: &ArgMatches) {
    migrations::find_migrations_directory()
        .unwrap_or_else(|_| {
            create_migrations_directory()
                .unwrap_or_else(handle_error)
        });

    database::setup_database(matches).unwrap_or_else(handle_error);
}

fn run_database_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("setup", Some(args)) => database::setup_database(args).unwrap_or_else(handle_error),
        ("reset", Some(args)) => database::reset_database(args).unwrap_or_else(handle_error),
        ("drop", Some(args)) => database::drop_database_command(args).unwrap_or_else(handle_error),
        _ => unreachable!("The cli parser should prevent reaching here"),
    };
}

/// Looks for a migrations directory in the current path and all parent paths,
/// and creates one in the same directory as the Cargo.toml if it can't find
/// one. It also sticks a .gitkeep in the directory so git will pick it up.
/// Returns a `DatabaseError::CargoTomlNotFound` if no Cargo.toml is found.
fn create_migrations_directory() -> DatabaseResult<PathBuf> {
    let project_root = try!(find_project_root());
    println!("Creating migrations/ directory at: {}", project_root
                                                        .join("migrations")
                                                        .display());
    try!(fs::create_dir(project_root.join("migrations")));
    try!(fs::File::create(project_root.join("migrations/.gitkeep")));
    Ok(project_root)
}

fn find_project_root() -> DatabaseResult<PathBuf> {
    search_for_cargo_toml_directory(&try!(env::current_dir()))
}

/// Searches for the directory that holds the project's Cargo.toml, and returns
/// the path if it found it, or returns a `DatabaseError::CargoTomlNotFound`.
fn search_for_cargo_toml_directory(path: &Path) -> DatabaseResult<PathBuf> {
    let toml_path = path.join("Cargo.toml");
    if toml_path.is_file() {
        Ok(path.to_owned())
    } else {
        path.parent().map(search_for_cargo_toml_directory)
            .unwrap_or(Err(DatabaseError::CargoTomlNotFound))
    }
}

/// Reverts the most recent migration, and then runs it again, all in a
/// transaction. If either part fails, the transaction is not committed.
fn redo_latest_migration<Conn>(conn: &Conn) where
        Conn: Connection,
        String: FromSql<VarChar, Conn::Backend>,
        for<'a> &'a NewMigration<'a>:
            Insertable<__diesel_schema_migrations::table, Conn::Backend>,
{
    conn.transaction(|| {
        let reverted_version = try!(migrations::revert_latest_migration(conn));
        migrations::run_migration_with_version(conn, &reverted_version, &mut stdout())
    }).unwrap_or_else(handle_error);
}

fn handle_error<E: Error, T>(error: E) -> T {
    panic!("{}", error);
}

// Converts an absolute path to a relative path, with the restriction that the
// target path must be in the same directory or above the current path.
fn convert_absolute_path_to_relative(target_path: &Path, mut current_path: &Path)
    -> PathBuf
{
    let mut result = PathBuf::new();

    while !target_path.starts_with(current_path) {
        result.push("..");
        match current_path.parent() {
            Some(parent) => current_path = parent,
            None => return target_path.into(),
        }
    }

    result.join(strip_prefix(target_path, current_path))
}

// FIXME: Remove all of this when 1.7 is stable
fn strip_prefix<'a>(target: &'a Path, base: &'a Path)
-> &'a Path {
    iter_after(target.components(), base.components())
        .map(|c| c.as_path()).unwrap()
}

fn iter_after<A, I, J>(mut iter: I, mut prefix: J) -> Option<I>
where I: Iterator<Item = A> + Clone,
      J: Iterator<Item = A>,
      A: PartialEq
{
    loop {
        let mut iter_next = iter.clone();
        match (iter_next.next(), prefix.next()) {
            (Some(x), Some(y)) => {
                if x != y {
                    return None;
                }
            }
            (Some(_), None) => return Some(iter),
            (None, None) => return Some(iter),
            (None, Some(_)) => return None,
        }
        iter = iter_next;
    }
}
// End FIXME

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use database_error::DatabaseError;

    use self::tempdir::TempDir;

    use std::fs;
    use std::path::PathBuf;

    use super::convert_absolute_path_to_relative;
    use super::search_for_cargo_toml_directory;

    #[test]
    fn toml_directory_find_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let toml_path = temp_path.join("Cargo.toml");

        fs::File::create(&toml_path).unwrap();

        assert_eq!(Ok(temp_path.clone()), search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn cargo_toml_not_found_if_no_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();

        assert_eq!(Err(DatabaseError::CargoTomlNotFound),
            search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn convert_absolute_path_to_relative_works() {
        assert_eq!(PathBuf::from("migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo")));
        assert_eq!(PathBuf::from("../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src")));
        assert_eq!(PathBuf::from("../../../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src/controllers/errors")));
        assert_eq!(PathBuf::from("12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations")));
        assert_eq!(PathBuf::from("../12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations/67890_create_post")));
    }
}
<commit_msg>Fix `diesel database setup` help, and specify subgroup help messages<commit_after>extern crate chrono;
extern crate clap;
extern crate diesel;
extern crate dotenv;

mod database_error;
#[macro_use]
mod database;

use chrono::*;
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};
#[cfg(feature = "postgres")]
use diesel::pg::PgConnection;
#[cfg(feature = "sqlite")]
use diesel::sqlite::SqliteConnection;
use diesel::migrations::schema::*;
use diesel::types::{FromSql, VarChar};
use diesel::{migrations, Connection, Insertable};
use std::error::Error;
use std::io::stdout;
use std::path::{PathBuf, Path};
use std::{env, fs};

use self::database_error::{DatabaseError, DatabaseResult};

fn main() {
    use self::dotenv::dotenv;
    dotenv().ok();

    let database_arg = Arg::with_name("DATABASE_URL")
        .long("database-url")
        .help("Specifies the database URL to connect to. Falls back to \
                   the DATABASE_URL environment variable if unspecified.")
        .global(true)
        .takes_value(true);

    let migration_subcommand = SubCommand::with_name("migration")
        .about("A group of commands for generating, running, and reverting \
                migrations.")
        .setting(AppSettings::VersionlessSubcommands)
        .arg(Arg::with_name("MIGRATION_DIRECTORY")
            .long("migration-dir")
            .help("The location of your migration directory. By default this \
                   will look for a directory called `migrations` in the \
                   current directory and its parents.")
            .takes_value(true)
            .global(true)
        ).subcommand(
            SubCommand::with_name("run")
                .about("Runs all pending migrations")
        ).subcommand(
            SubCommand::with_name("revert")
                .about("Reverts the latest run migration")
        ).subcommand(
            SubCommand::with_name("redo")
                .about("Reverts and re-runs the latest migration. Useful \
                      for testing that a migration can in fact be reverted.")
        ).subcommand(
            SubCommand::with_name("generate")
                .about("Generate a new migration with the given name, and \
                      the current timestamp as the version")
                .arg(Arg::with_name("MIGRATION_NAME")
                     .help("The name of the migration to create")
                     .required(true)
                 )
                .arg(Arg::with_name("MIGRATION_VERSION")
                     .long("version")
                     .help("The version number to use when generating the migration. \
                            Defaults to the current timestamp, which should suffice \
                            for most use cases.")
                     .takes_value(true)
                )
        ).setting(AppSettings::SubcommandRequiredElseHelp);

    let setup_subcommand = SubCommand::with_name("setup")
        .about("Creates the migrations directory, creates the database \
                specified in your DATABASE_URL, and runs existing migrations.");

    let database_subcommand = SubCommand::with_name("database")
        .about("A group of commands for setting up and resetting your database.")
        .setting(AppSettings::VersionlessSubcommands)
        .subcommand(
            SubCommand::with_name("setup")
                .about("Creates the database specified in your DATABASE_URL, \
                        and then runs any existing migrations.")
        ).subcommand(
            SubCommand::with_name("reset")
                .about("Resets your database by dropping the database specified \
                        in your DATABASE_URL and then running `diesel database setup`.")
        ).subcommand(
            SubCommand::with_name("drop")
                .about("Drops the database specified in your DATABASE_URL.")
                .setting(AppSettings::Hidden)
        ).setting(AppSettings::SubcommandRequiredElseHelp);

    let matches = App::new("diesel")
        .version(env!("CARGO_PKG_VERSION"))
        .setting(AppSettings::VersionlessSubcommands)
        .after_help("You can also run `diesel SUBCOMMAND -h` to get more information about that subcommand.")
        .arg(database_arg)
        .subcommand(migration_subcommand)
        .subcommand(setup_subcommand)
        .subcommand(database_subcommand)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .get_matches();

    match matches.subcommand() {
        ("migration", Some(matches)) => run_migration_command(matches),
        ("setup", Some(matches)) => run_setup_command(matches),
        ("database", Some(matches)) => run_database_command(matches),
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

fn run_migration_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("run", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, migrations::run_pending_migrations)
                .unwrap_or_else(handle_error);
        }
        ("revert", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, migrations::revert_latest_migration)
                .unwrap_or_else(handle_error);
        }
        ("redo", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, redo_latest_migration);
        }
        ("generate", Some(args)) => {
            let migration_name = args.value_of("MIGRATION_NAME").unwrap();
            let version = migration_version(args);
            let versioned_name = format!("{}_{}", version, migration_name);
            let migration_dir = migrations_dir(args).join(versioned_name);
            fs::create_dir(&migration_dir).unwrap();

            let migration_dir_relative = convert_absolute_path_to_relative(
                &migration_dir,
                &env::current_dir().unwrap()
            );

            let up_path = migration_dir.join("up.sql");
            println!("Creating {}", migration_dir_relative.join("up.sql").display());
            fs::File::create(up_path).unwrap();
            let down_path = migration_dir.join("down.sql");
            println!("Creating {}", migration_dir_relative.join("down.sql").display());
            fs::File::create(down_path).unwrap();
        }
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

use std::fmt::Display;
fn migration_version<'a>(matches: &'a ArgMatches) -> Box<Display + 'a> {
    matches.value_of("MIGRATION_VERSION").map(|s| Box::new(s) as Box<Display>)
        .unwrap_or_else(|| Box::new(Local::now().format("%Y%m%d%H%M%S")))
}

fn migrations_dir(matches: &ArgMatches) -> PathBuf {
    matches.value_of("MIGRATION_DIRECTORY")
        .map(PathBuf::from)
        .or_else(|| {
            env::var("MIGRATION_DIRECTORY").map(PathBuf::from).ok()
        }).unwrap_or_else(|| {
            migrations::find_migrations_directory()
                .unwrap_or_else(handle_error)
        })
}

fn run_setup_command(matches: &ArgMatches) {
    migrations::find_migrations_directory()
        .unwrap_or_else(|_| {
            create_migrations_directory()
                .unwrap_or_else(handle_error)
        });

    database::setup_database(matches).unwrap_or_else(handle_error);
}

fn run_database_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("setup", Some(args)) => database::setup_database(args).unwrap_or_else(handle_error),
        ("reset", Some(args)) => database::reset_database(args).unwrap_or_else(handle_error),
        ("drop", Some(args)) => database::drop_database_command(args).unwrap_or_else(handle_error),
        _ => unreachable!("The cli parser should prevent reaching here"),
    };
}

/// Looks for a migrations directory in the current path and all parent paths,
/// and creates one in the same directory as the Cargo.toml if it can't find
/// one. It also sticks a .gitkeep in the directory so git will pick it up.
/// Returns a `DatabaseError::CargoTomlNotFound` if no Cargo.toml is found.
fn create_migrations_directory() -> DatabaseResult<PathBuf> {
    let project_root = try!(find_project_root());
    println!("Creating migrations/ directory at: {}", project_root
                                                        .join("migrations")
                                                        .display());
    try!(fs::create_dir(project_root.join("migrations")));
    try!(fs::File::create(project_root.join("migrations/.gitkeep")));
    Ok(project_root)
}

fn find_project_root() -> DatabaseResult<PathBuf> {
    search_for_cargo_toml_directory(&try!(env::current_dir()))
}

/// Searches for the directory that holds the project's Cargo.toml, and returns
/// the path if it found it, or returns a `DatabaseError::CargoTomlNotFound`.
fn search_for_cargo_toml_directory(path: &Path) -> DatabaseResult<PathBuf> {
    let toml_path = path.join("Cargo.toml");
    if toml_path.is_file() {
        Ok(path.to_owned())
    } else {
        path.parent().map(search_for_cargo_toml_directory)
            .unwrap_or(Err(DatabaseError::CargoTomlNotFound))
    }
}

/// Reverts the most recent migration, and then runs it again, all in a
/// transaction. If either part fails, the transaction is not committed.
fn redo_latest_migration<Conn>(conn: &Conn) where
        Conn: Connection,
        String: FromSql<VarChar, Conn::Backend>,
        for<'a> &'a NewMigration<'a>:
            Insertable<__diesel_schema_migrations::table, Conn::Backend>,
{
    conn.transaction(|| {
        let reverted_version = try!(migrations::revert_latest_migration(conn));
        migrations::run_migration_with_version(conn, &reverted_version, &mut stdout())
    }).unwrap_or_else(handle_error);
}

fn handle_error<E: Error, T>(error: E) -> T {
    panic!("{}", error);
}

// Converts an absolute path to a relative path, with the restriction that the
// target path must be in the same directory or above the current path.
fn convert_absolute_path_to_relative(target_path: &Path, mut current_path: &Path)
    -> PathBuf
{
    let mut result = PathBuf::new();

    while !target_path.starts_with(current_path) {
        result.push("..");
        match current_path.parent() {
            Some(parent) => current_path = parent,
            None => return target_path.into(),
        }
    }

    result.join(strip_prefix(target_path, current_path))
}

// FIXME: Remove all of this when 1.7 is stable
fn strip_prefix<'a>(target: &'a Path, base: &'a Path)
-> &'a Path {
    iter_after(target.components(), base.components())
        .map(|c| c.as_path()).unwrap()
}

fn iter_after<A, I, J>(mut iter: I, mut prefix: J) -> Option<I>
where I: Iterator<Item = A> + Clone,
      J: Iterator<Item = A>,
      A: PartialEq
{
    loop {
        let mut iter_next = iter.clone();
        match (iter_next.next(), prefix.next()) {
            (Some(x), Some(y)) => {
                if x != y {
                    return None;
                }
            }
            (Some(_), None) => return Some(iter),
            (None, None) => return Some(iter),
            (None, Some(_)) => return None,
        }
        iter = iter_next;
    }
}
// End FIXME

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use database_error::DatabaseError;

    use self::tempdir::TempDir;

    use std::fs;
    use std::path::PathBuf;

    use super::convert_absolute_path_to_relative;
    use super::search_for_cargo_toml_directory;

    #[test]
    fn toml_directory_find_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let toml_path = temp_path.join("Cargo.toml");

        fs::File::create(&toml_path).unwrap();

        assert_eq!(Ok(temp_path.clone()), search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn cargo_toml_not_found_if_no_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();

        assert_eq!(Err(DatabaseError::CargoTomlNotFound),
            search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn convert_absolute_path_to_relative_works() {
        assert_eq!(PathBuf::from("migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo")));
        assert_eq!(PathBuf::from("../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src")));
        assert_eq!(PathBuf::from("../../../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src/controllers/errors")));
        assert_eq!(PathBuf::from("12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations")));
        assert_eq!(PathBuf::from("../12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations/67890_create_post")));
    }
}
<|endoftext|>


<commit_before>use wasm_bindgen::prelude::*;

use anyhow::Error;
use lancelot::{
    analysis::dis::{self, zydis},
    arch::Arch,
    aspace::AddressSpace,
    loader::pe::{PEError, PE as lPE},
    module::ModuleError,
    pagemap::PageMapError,
    util::UtilError,
};

#[cfg(feature = "wee_alloc")]
#[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;

#[wasm_bindgen]
pub fn version() -> String {
    String::from("v0.7.0")
}

#[wasm_bindgen(start)]
pub fn main_js() -> Result<(), JsValue> {
    #[cfg(debug_assertions)]
    console_error_panic_hook::set_once();

    Ok(())
}

// for better function signatures
// no functional difference.
type JsError = JsValue;

fn to_value_error(e: anyhow::Error) -> JsError {
    js_sys::Error::new(&format!("{}", e)).into()
}

fn to_js_err(e: Error) -> JsError {
    match e.downcast_ref::<PEError>() {
        Some(PEError::FormatNotSupported(_)) => return to_value_error(e),
        Some(PEError::MalformedPEFile(_)) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<ModuleError>() {
        Some(ModuleError::InvalidAddress(_)) => return to_value_error(e),
        None => (),
    };

    match e.downcast_ref::<UtilError>() {
        Some(UtilError::FileAccess) => return to_value_error(e),
        Some(UtilError::FileFormat) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<PageMapError>() {
        Some(PageMapError::NotMapped) => return to_value_error(e),
        None => (),
    };

    to_value_error(e)
}

#[wasm_bindgen]
pub fn from_bytes(buf: Vec<u8>) -> Result<PE, JsError> {
    let pe = lPE::from_bytes(&buf).map_err(to_js_err)?;
    let dec = dis::get_disassembler(&pe.module).map_err(to_js_err)?;
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("valid formatter options");
    Ok(PE {
        inner: pe,
        decoder: dec,
        formatter,
    })
}

#[wasm_bindgen]
#[derive(Clone, Copy)]
pub struct AddressRange {
    pub start: u64,
    pub end:   u64,
}

#[wasm_bindgen]
impl AddressRange {
    #[wasm_bindgen(getter)]
    pub fn size(&self) -> u64 {
        self.end - self.start
    }
}

#[wasm_bindgen]
pub struct Section {
    pub physical_range: AddressRange,
    pub virtual_range:  AddressRange,
    // String are not clone, so we have to manually implement getters
    permissions:        String,
    name:               String,
}

#[wasm_bindgen]
impl Section {
    #[wasm_bindgen(getter)]
    pub fn permissions(&self) -> String {
        self.permissions.to_string()
    }

    #[wasm_bindgen(getter)]
    pub fn name(&self) -> String {
        self.name.clone()
    }
}

impl From<&lancelot::module::Section> for Section {
    fn from(other: &lancelot::module::Section) -> Section {
        Section {
            physical_range: AddressRange {
                start: other.physical_range.start,
                end:   other.physical_range.end,
            },
            virtual_range:  AddressRange {
                start: other.virtual_range.start,
                end:   other.virtual_range.end,
            },
            permissions:    other.permissions.to_string(),
            name:           other.name.to_string(),
        }
    }
}

#[wasm_bindgen]
pub struct Instruction {
    pub address: u64,
    pub size:    u8,
    bytes:       Vec<u8>,
    string:      String,
}

#[wasm_bindgen]
impl Instruction {
    #[wasm_bindgen(getter)]
    pub fn bytes(&self) -> Vec<u8> {
        self.bytes.clone()
    }

    #[wasm_bindgen(getter)]
    pub fn string(&self) -> String {
        self.string.to_string()
    }
}

#[allow(clippy::upper_case_acronyms)]
enum StringEncoding {
    ASCII,
    UTF16,
}

#[wasm_bindgen(js_name = "String")]
pub struct EncodedString {
    encoding:    StringEncoding,
    pub address: u64,
    pub size:    usize,
    string:      String,
}

#[wasm_bindgen(js_class = "String")]
impl EncodedString {
    #[wasm_bindgen(getter, typescript_type = "string")]
    pub fn encoding(&self) -> JsValue {
        match self.encoding {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            StringEncoding::ASCII => JsValue::from(wasm_bindgen::intern("ASCII")),
            StringEncoding::UTF16 => JsValue::from(wasm_bindgen::intern("UTF-16")),
        }
    }

    #[wasm_bindgen(getter)]
    pub fn string(&self) -> String {
        self.string.to_string()
    }
}

#[wasm_bindgen]
pub struct PE {
    inner:     lPE,
    decoder:   zydis::Decoder,
    formatter: zydis::Formatter,
}

#[wasm_bindgen]
impl PE {
    #[wasm_bindgen(getter, typescript_type = "string")]
    pub fn arch(&self) -> JsValue {
        match self.inner.module.arch {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            Arch::X32 => JsValue::from(wasm_bindgen::intern("x32")),
            Arch::X64 => JsValue::from(wasm_bindgen::intern("x64")),
        }
    }

    #[wasm_bindgen(getter)]
    pub fn base_address(&self) -> u64 {
        self.inner.module.address_space.base_address
    }

    // not cheap, so its a function, not getter
    #[wasm_bindgen]
    pub fn functions(&self) -> Result<Vec<u64>, JsError> {
        Ok(lancelot::analysis::pe::find_functions(&self.inner)
            .map_err(to_js_err)?
            .into_iter()
            .filter(|f| matches!(f, lancelot::analysis::pe::Function::Local(_)))
            .map(|f| match f {
                lancelot::analysis::pe::Function::Local(va) => va,
                _ => unreachable!(),
            })
            .collect())
    }

    // Vec<T> cannot be serialized by wasm-bindgen
    // so we have to manually convert to Vec<JsValue>
    // and then annotate the TS type.
    //
    // NB: JS now owns the objects, must explicitly drop.
    #[wasm_bindgen(getter, typescript_type = "Array<Section>")]
    pub fn sections(&self) -> Vec<JsValue> {
        self.inner
            .module
            .sections
            .iter()
            .map(Section::from)
            .map(JsValue::from)
            .collect()
    }

    #[wasm_bindgen]
    pub fn read_bytes(&self, va: u64, size: usize) -> Result<Vec<u8>, JsError> {
        self.inner.module.address_space.read_bytes(va, size).map_err(to_js_err)
    }

    #[wasm_bindgen]
    pub fn read_insn(&self, va: u64) -> Result<Instruction, JsError> {
        let mut insn_buf = [0u8; 16];
        self.inner
            .module
            .address_space
            .read_into(va, &mut insn_buf)
            .map_err(to_js_err)?;

        if let Ok(Some(insn)) = self.decoder.decode(&insn_buf) {
            let mut out_buf = [0u8; 200];
            let mut out_buf = zydis::OutputBuffer::new(&mut out_buf[..]);
            self.formatter
                .format_instruction(&insn, &mut out_buf, Some(va), None)
                .map_err(|e| -> JsError {
                    js_sys::Error::new(&format!("failed to format instruction: {}", e)).into()
                })?;

            Ok(Instruction {
                address: va,
                size:    insn.length,
                bytes:   insn_buf[..insn.length as usize].to_vec(),
                string:  out_buf
                    .as_str()
                    .expect("disassembly should always be utf-8")
                    .to_string(),
            })
        } else {
            Err(js_sys::Error::new("invalid instruction").into())
        }
    }

    #[wasm_bindgen(typescript_type = "Array<EncodedString>")]
    pub fn strings(&self) -> Vec<JsValue> {
        let mut ret: Vec<EncodedString> = Default::default();

        for (range, s) in lancelot::util::find_ascii_strings(&self.inner.buf) {
            if let Ok(va) = self.inner.module.virtual_address(range.start as u64) {
                ret.push(EncodedString {
                    encoding: StringEncoding::ASCII,
                    address:  va,
                    size:     range.end - range.start,
                    string:   s,
                });
            }
        }

        for (range, s) in lancelot::util::find_unicode_strings(&self.inner.buf) {
            if let Ok(va) = self.inner.module.virtual_address(range.start as u64) {
                ret.push(EncodedString {
                    encoding: StringEncoding::UTF16,
                    address:  va,
                    size:     range.end - range.start,
                    string:   s,
                });
            }
        }

        ret.sort_by_key(|string| string.address);

        ret.into_iter().map(JsValue::from).collect()
    }
}
<commit_msg>jslancelot: expose layout and callgraph<commit_after>use std::collections::BTreeMap;

use wasm_bindgen::prelude::*;

use anyhow::Error;
use lancelot::{
    analysis::dis::{self, zydis},
    arch::Arch,
    aspace::AddressSpace,
    loader::pe::{PEError, PE as lPE},
    module::ModuleError,
    pagemap::PageMapError,
    util::UtilError,
};

#[cfg(feature = "wee_alloc")]
#[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;

#[wasm_bindgen]
pub fn version() -> String {
    String::from("v0.7.0")
}

#[wasm_bindgen(start)]
pub fn main_js() -> Result<(), JsValue> {
    #[cfg(debug_assertions)]
    console_error_panic_hook::set_once();

    Ok(())
}

// for better function signatures
// no functional difference.
type JsError = JsValue;

fn to_value_error(e: anyhow::Error) -> JsError {
    js_sys::Error::new(&format!("{}", e)).into()
}

fn to_js_err(e: Error) -> JsError {
    match e.downcast_ref::<PEError>() {
        Some(PEError::FormatNotSupported(_)) => return to_value_error(e),
        Some(PEError::MalformedPEFile(_)) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<ModuleError>() {
        Some(ModuleError::InvalidAddress(_)) => return to_value_error(e),
        None => (),
    };

    match e.downcast_ref::<UtilError>() {
        Some(UtilError::FileAccess) => return to_value_error(e),
        Some(UtilError::FileFormat) => return to_value_error(e),
        None => (),
    };

    #[allow(clippy::single_match)]
    match e.downcast_ref::<PageMapError>() {
        Some(PageMapError::NotMapped) => return to_value_error(e),
        None => (),
    };

    to_value_error(e)
}

#[wasm_bindgen]
pub fn from_bytes(buf: Vec<u8>) -> Result<PE, JsError> {
    let pe = lPE::from_bytes(&buf).map_err(to_js_err)?;
    let dec = dis::get_disassembler(&pe.module).map_err(to_js_err)?;
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("valid formatter options");
    Ok(PE {
        inner: pe,
        decoder: dec,
        formatter,
    })
}

#[wasm_bindgen]
#[derive(Clone, Copy)]
pub struct AddressRange {
    pub start: u64,
    pub end:   u64,
}

#[wasm_bindgen]
impl AddressRange {
    #[wasm_bindgen(getter)]
    pub fn size(&self) -> u64 {
        self.end - self.start
    }
}

#[wasm_bindgen]
pub struct Section {
    pub physical_range: AddressRange,
    pub virtual_range:  AddressRange,
    // String are not clone, so we have to manually implement getters
    permissions:        String,
    name:               String,
}

#[wasm_bindgen]
impl Section {
    #[wasm_bindgen(getter)]
    pub fn permissions(&self) -> String {
        self.permissions.to_string()
    }

    #[wasm_bindgen(getter)]
    pub fn name(&self) -> String {
        self.name.clone()
    }
}

impl From<&lancelot::module::Section> for Section {
    fn from(other: &lancelot::module::Section) -> Section {
        Section {
            physical_range: AddressRange {
                start: other.physical_range.start,
                end:   other.physical_range.end,
            },
            virtual_range:  AddressRange {
                start: other.virtual_range.start,
                end:   other.virtual_range.end,
            },
            permissions:    other.permissions.to_string(),
            name:           other.name.to_string(),
        }
    }
}

#[wasm_bindgen]
pub struct Instruction {
    pub address: u64,
    pub size:    u8,
    bytes:       Vec<u8>,
    string:      String,
}

#[wasm_bindgen]
impl Instruction {
    #[wasm_bindgen(getter)]
    pub fn bytes(&self) -> Vec<u8> {
        self.bytes.clone()
    }

    #[wasm_bindgen(getter)]
    pub fn string(&self) -> String {
        self.string.to_string()
    }
}

#[allow(clippy::upper_case_acronyms)]
enum StringEncoding {
    ASCII,
    UTF16,
}

#[wasm_bindgen(js_name = "String")]
pub struct EncodedString {
    encoding:    StringEncoding,
    pub address: u64,
    pub size:    usize,
    string:      String,
}

#[wasm_bindgen(js_class = "String")]
impl EncodedString {
    #[wasm_bindgen(getter, typescript_type = "string")]
    pub fn encoding(&self) -> JsValue {
        match self.encoding {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            StringEncoding::ASCII => JsValue::from(wasm_bindgen::intern("ASCII")),
            StringEncoding::UTF16 => JsValue::from(wasm_bindgen::intern("UTF-16")),
        }
    }

    #[wasm_bindgen(getter)]
    pub fn string(&self) -> String {
        self.string.to_string()
    }
}

#[wasm_bindgen]
#[derive(Clone)]
pub struct Flow {
    inner: lancelot::analysis::cfg::Flow,
    /* type: String,
     * target: u64, */
}

#[wasm_bindgen]
impl Flow {
    #[wasm_bindgen(getter, js_name = "type", typescript_type = "string")]
    pub fn type_(&self) -> JsValue {
        match self.inner {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            lancelot::analysis::cfg::Flow::Call(_) => JsValue::from(wasm_bindgen::intern("call")),
            lancelot::analysis::cfg::Flow::Fallthrough(_) => JsValue::from(wasm_bindgen::intern("fallthrough")),
            lancelot::analysis::cfg::Flow::UnconditionalJump(_) => {
                JsValue::from(wasm_bindgen::intern("unconditional jump"))
            }
            lancelot::analysis::cfg::Flow::ConditionalJump(_) => {
                JsValue::from(wasm_bindgen::intern("conditional jump"))
            }
            lancelot::analysis::cfg::Flow::ConditionalMove(_) => {
                JsValue::from(wasm_bindgen::intern("conditional move"))
            }
        }
    }

    #[wasm_bindgen(getter)]
    pub fn target(&self) -> u64 {
        self.inner.va()
    }
}

impl From<lancelot::analysis::cfg::Flow> for Flow {
    fn from(other: lancelot::analysis::cfg::Flow) -> Flow {
        Flow { inner: other }
    }
}

#[wasm_bindgen]
#[derive(Clone)]
pub struct BasicBlock {
    pub address: u64,
    pub size:    u64,

    successors: Vec<lancelot::analysis::cfg::Flow>,

    instructions: Vec<u64>,
}

#[wasm_bindgen]
impl BasicBlock {
    // Vec<Flow> cannot be serialized by wasm-bindgen
    // so we have to manually convert to Vec<JsValue>
    // and then annotate the TS type.
    //
    // NB: JS now owns the objects, must explicitly drop.
    #[wasm_bindgen(getter, typescript_type = "Array<Flow>")]
    pub fn successors(&self) -> Vec<JsValue> {
        self.successors
            .iter()
            .cloned()
            .map(Flow::from)
            .map(JsValue::from)
            .collect()
    }

    // Vec<u64> is not Copy so cannot be serialized by wasm-bindgen.
    #[wasm_bindgen(getter)]
    pub fn instructions(&self) -> Vec<u64> {
        self.instructions.clone()
    }
}

#[wasm_bindgen]
#[derive(Clone)]
pub struct Function {
    pub address:  u64,
    basic_blocks: Vec<BasicBlock>,
}

#[wasm_bindgen]
impl Function {
    // Vec<BasicBlock> cannot be serialized by wasm-bindgen
    // so we have to manually convert to Vec<JsValue>
    // and then annotate the TS type.
    //
    // NB: JS now owns the objects, must explicitly drop.
    #[wasm_bindgen(getter, typescript_type = "Array<BasicBlock>")]
    pub fn basic_blocks(&self) -> Vec<JsValue> {
        self.basic_blocks.iter().cloned().map(JsValue::from).collect()
    }
}

// create a js Map<bigint, bigint[]> from a rust BTreeMap<u64, Vec<u64>>
fn jsvalue_from_addresses_by_address(other: &BTreeMap<u64, Vec<u64>>) -> JsValue {
    let ret = js_sys::Map::new();

    for (&k, v) in other.iter() {
        let l = js_sys::Array::new_with_length(v.len() as u32);
        for (i, &vv) in v.iter().enumerate() {
            l.set(i as u32, JsValue::from(vv));
        }

        ret.set(&JsValue::from(k), &JsValue::from(l));
    }

    ret.into()
}

#[wasm_bindgen]
#[derive(Clone)]
pub struct CallGraph {
    calls_to:                   JsValue,
    calls_from:                 JsValue,
    function_call_instructions: JsValue,
    call_instruction_functions: JsValue,
}

#[wasm_bindgen]
impl CallGraph {
    #[wasm_bindgen(getter, typescript_type = "Map<bigint, bigint[]>")]
    pub fn calls_to(&self) -> JsValue {
        self.calls_to.clone()
    }

    #[wasm_bindgen(getter, typescript_type = "Map<bigint, bigint[]>")]
    pub fn calls_from(&self) -> JsValue {
        self.calls_from.clone()
    }

    #[wasm_bindgen(getter, typescript_type = "Map<bigint, bigint[]>")]
    pub fn function_call_instructions(&self) -> JsValue {
        self.function_call_instructions.clone()
    }

    #[wasm_bindgen(getter, typescript_type = "Map<bigint, bigint[]>")]
    pub fn call_instruction_functions(&self) -> JsValue {
        self.call_instruction_functions.clone()
    }
}

#[wasm_bindgen]
pub struct Layout {
    functions:  Vec<Function>,
    call_graph: CallGraph,
}

#[wasm_bindgen]
impl Layout {
    #[wasm_bindgen(getter, typescript_type = "Map<bigint, Function>")]
    pub fn functions(&self) -> JsValue {
        let ret = js_sys::Map::new();

        for f in self.functions.iter().cloned() {
            ret.set(&JsValue::from(f.address), &JsValue::from(f));
        }

        ret.into()
    }

    #[wasm_bindgen(getter, typescript_type = "CallGraph")]
    pub fn call_graph(&self) -> JsValue {
        JsValue::from(self.call_graph.clone())
    }
}

#[wasm_bindgen]
pub struct PE {
    inner:     lPE,
    decoder:   zydis::Decoder,
    formatter: zydis::Formatter,
}

#[wasm_bindgen]
impl PE {
    #[wasm_bindgen(getter, typescript_type = "string")]
    pub fn arch(&self) -> JsValue {
        match self.inner.module.arch {
            // we expect these strings may be returned many times,
            // so we intern them for perf
            // https://docs.rs/wasm-bindgen/0.2.78/wasm_bindgen/fn.intern.html
            Arch::X32 => JsValue::from(wasm_bindgen::intern("x32")),
            Arch::X64 => JsValue::from(wasm_bindgen::intern("x64")),
        }
    }

    #[wasm_bindgen(getter)]
    pub fn base_address(&self) -> u64 {
        self.inner.module.address_space.base_address
    }

    // not cheap, so its a function, not getter
    #[wasm_bindgen]
    pub fn functions(&self) -> Result<Vec<u64>, JsError> {
        Ok(lancelot::analysis::pe::find_functions(&self.inner)
            .map_err(to_js_err)?
            .into_iter()
            .filter(|f| matches!(f, lancelot::analysis::pe::Function::Local(_)))
            .map(|f| match f {
                lancelot::analysis::pe::Function::Local(va) => va,
                _ => unreachable!(),
            })
            .collect())
    }

    // Vec<T> cannot be serialized by wasm-bindgen
    // so we have to manually convert to Vec<JsValue>
    // and then annotate the TS type.
    //
    // NB: JS now owns the objects, must explicitly drop.
    #[wasm_bindgen(getter, typescript_type = "Array<Section>")]
    pub fn sections(&self) -> Vec<JsValue> {
        self.inner
            .module
            .sections
            .iter()
            .map(Section::from)
            .map(JsValue::from)
            .collect()
    }

    #[wasm_bindgen]
    pub fn read_bytes(&self, va: u64, size: usize) -> Result<Vec<u8>, JsError> {
        self.inner.module.address_space.read_bytes(va, size).map_err(to_js_err)
    }

    #[wasm_bindgen]
    pub fn read_insn(&self, va: u64) -> Result<Instruction, JsError> {
        let mut insn_buf = [0u8; 16];
        self.inner
            .module
            .address_space
            .read_into(va, &mut insn_buf)
            .map_err(to_js_err)?;

        if let Ok(Some(insn)) = self.decoder.decode(&insn_buf) {
            let mut out_buf = [0u8; 200];
            let mut out_buf = zydis::OutputBuffer::new(&mut out_buf[..]);
            self.formatter
                .format_instruction(&insn, &mut out_buf, Some(va), None)
                .map_err(|e| -> JsError {
                    js_sys::Error::new(&format!("failed to format instruction: {}", e)).into()
                })?;

            Ok(Instruction {
                address: va,
                size:    insn.length,
                bytes:   insn_buf[..insn.length as usize].to_vec(),
                string:  out_buf
                    .as_str()
                    .expect("disassembly should always be utf-8")
                    .to_string(),
            })
        } else {
            Err(js_sys::Error::new("invalid instruction").into())
        }
    }

    #[wasm_bindgen(typescript_type = "Array<EncodedString>")]
    pub fn strings(&self) -> Vec<JsValue> {
        let mut ret: Vec<EncodedString> = Default::default();

        for (range, s) in lancelot::util::find_ascii_strings(&self.inner.buf) {
            if let Ok(va) = self.inner.module.virtual_address(range.start as u64) {
                ret.push(EncodedString {
                    encoding: StringEncoding::ASCII,
                    address:  va,
                    size:     range.end - range.start,
                    string:   s,
                });
            }
        }

        for (range, s) in lancelot::util::find_unicode_strings(&self.inner.buf) {
            if let Ok(va) = self.inner.module.virtual_address(range.start as u64) {
                ret.push(EncodedString {
                    encoding: StringEncoding::UTF16,
                    address:  va,
                    size:     range.end - range.start,
                    string:   s,
                });
            }
        }

        ret.sort_by_key(|string| string.address);

        ret.into_iter().map(JsValue::from).collect()
    }

    #[wasm_bindgen]
    pub fn layout(&self) -> Result<Layout, JsValue> {
        use lancelot::{analysis::cfg::CFG, aspace::AddressSpace, VA};
        use std::collections::{BTreeMap, BTreeSet};

        let mut cfgs: BTreeMap<VA, CFG> = Default::default();
        for &function in lancelot::analysis::pe::find_function_starts(&self.inner)
            .map_err(to_js_err)?
            .iter()
        {
            if let Ok(cfg) = lancelot::analysis::cfg::build_cfg(&self.inner.module, function) {
                cfgs.insert(function, cfg);
            }
        }

        let imports: BTreeSet<VA> = lancelot::analysis::pe::get_imports(&self.inner)
            .map_err(to_js_err)?
            .keys()
            .cloned()
            .collect();
        let cg =
            lancelot::analysis::call_graph::build_call_graph(&self.inner.module, &cfgs, &imports).map_err(to_js_err)?;

        let mut functions = Vec::with_capacity(cfgs.len());

        for (&fva, cfg) in cfgs.iter() {
            let mut basic_blocks = Vec::with_capacity(cfg.basic_blocks.len());

            for (&bbva, bb) in cfg.basic_blocks.iter() {
                // estimate each instruction is 2 bytes, which is probably too conservative,
                // but enough to hint the allocation here.
                //
                // supported by here: https://ieeexplore.ieee.org/document/5645851
                // > The results show that the average instruction size length is about 2 bytes.
                let mut instructions = Vec::with_capacity((bb.length / 2) as usize);

                let buf = self
                    .inner
                    .module
                    .address_space
                    .read_bytes(bb.address, bb.length as usize)
                    .map_err(to_js_err)?;
                for (offset, insn) in lancelot::analysis::dis::linear_disassemble(&self.decoder, &buf) {
                    if let Ok(Some(_)) = insn {
                        let insnva = bb.address + offset as u64;
                        instructions.push(insnva);
                    }
                }

                basic_blocks.push(BasicBlock {
                    address: bbva,
                    size: bb.length as u64,
                    successors: bb.successors.to_vec(),
                    instructions,
                })
            }

            functions.push(Function {
                address: fva,
                basic_blocks,
            });
        }

        Ok(Layout {
            functions,
            call_graph: CallGraph {
                calls_to:                   jsvalue_from_addresses_by_address(&cg.calls_to),
                calls_from:                 jsvalue_from_addresses_by_address(&cg.calls_from),
                function_call_instructions: jsvalue_from_addresses_by_address(&cg.function_call_instructions),
                call_instruction_functions: jsvalue_from_addresses_by_address(&cg.call_instruction_functions),
            },
        })
    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use super::Configuration;
use ::client::Context;
use ::model::Message;
use ::model::Permissions;
use std::collections::HashMap;

pub type Check = Fn(&mut Context, &Message) -> bool + Send + Sync + 'static;
pub type Exec = Fn(&mut Context, &Message, Vec<String>) -> Result<(), String> + Send + Sync + 'static;
pub type Help = Fn(&mut Context, &Message, HashMap<String, Arc<CommandGroup>>, &[String]) -> Result<(), String> + Send + Sync + 'static;
pub type BeforeHook = Fn(&mut Context, &Message, &String) -> bool + Send + Sync + 'static;
pub type AfterHook = Fn(&mut Context, &Message, &String, Result<(), String>) + Send + Sync + 'static;
#[doc(hidden)]
pub type InternalCommand = Arc<Command>;
pub type PrefixCheck = Fn(&mut Context) -> Option<String> + Send + Sync + 'static;

#[doc(hidden)]
pub enum CommandOrAlias {
    Alias(String),
    Command(InternalCommand),
}

/// Command function type. Allows to access internal framework things inside
/// your commands.
pub enum CommandType {
    StringResponse(String),
    Basic(Box<Exec>),
    WithCommands(Box<Help>),
}

#[derive(Default)]
pub struct CommandGroup {
    pub prefix: Option<String>,
    pub commands: HashMap<String, CommandOrAlias>,
}

/// Command struct used to store commands internally.
pub struct Command {
    /// A set of checks to be called prior to executing the command. The checks
    /// will short-circuit on the first check that returns `false`.
    pub checks: Vec<Box<Check>>,
    /// Function called when the command is called.
    pub exec: CommandType,
    /// Ratelimit bucket.
    pub bucket: Option<String>,
    /// Command description, used by other commands.
    pub desc: Option<String>,
    /// Example arguments, used by other commands.
    pub example: Option<String>,
    /// Command usage schema, used by other commands.
    pub usage: Option<String>,
    /// Whether arguments should be parsed using quote parser or not.
    pub use_quotes: bool,
    /// Minumum amount of arguments that should be passed.
    pub min_args: Option<i32>,
    /// Maximum amount of arguments that can be passed.
    pub max_args: Option<i32>,
    /// Permissions required to use this command.
    pub required_permissions: Permissions,
    /// Whether command should be displayed in help list or not, used by other commands.
    pub help_available: bool,
    /// Whether command can be used only privately or not.
    pub dm_only: bool,
    /// Whether command can be used only in guilds or not.
    pub guild_only: bool,
    /// Whether command can only be used by owners or not.
    pub owners_only: bool,
    #[doc(hidden)]
    pub aliases: Vec<String>,
}

impl Command {
    pub fn new<F>(f: F) -> Self
        where F: Fn(&mut Context, &Message, Vec<String>) -> Result<(), String> + Send + Sync + 'static {
        Command {
            aliases: Vec::new(),
            checks: Vec::default(),
            exec: CommandType::Basic(Box::new(f)),
            desc: None,
            usage: None,
            example: None,
            use_quotes: false,
            dm_only: false,
            bucket: None,
            guild_only: false,
            help_available: true,
            min_args: None,
            max_args: None,
            owners_only: false,
            required_permissions: Permissions::empty(),
        }
    }
}

pub fn positions(ctx: &mut Context, content: &str, conf: &Configuration) -> Option<Vec<usize>> {
    if !conf.prefixes.is_empty() || conf.dynamic_prefix.is_some() {
        // Find out if they were mentioned. If not, determine if the prefix
        // was used. If not, return None.
        let mut positions: Vec<usize> = vec![];

        if let Some(mention_end) = find_mention_end(content, conf) {
            positions.push(mention_end);
        } else if let Some(ref func) = conf.dynamic_prefix {
            if let Some(x) = func(ctx) {
                positions.push(x.len());
            } else {
                for n in conf.prefixes.clone() {
                    if content.starts_with(&n) {
                        positions.push(n.len());
                    }
                }
            }
        } else {
            for n in conf.prefixes.clone() {
                if content.starts_with(&n) {
                    positions.push(n.len());
                }
            }
        };

        if positions.is_empty() {
            return None;
        }

        if conf.allow_whitespace {
            let pos = *unsafe { positions.get_unchecked(0) };

            positions.insert(0, pos + 1);
        }

        Some(positions)
    } else if conf.on_mention.is_some() {
        match find_mention_end(content, conf) {
            Some(mention_end) => {
                let mut positions = vec![mention_end];

                if conf.allow_whitespace {
                    positions.insert(0, mention_end + 1);
                }

                Some(positions)
            },
            None => None,
        }
    } else {
        None
    }
}

fn find_mention_end(content: &str, conf: &Configuration) -> Option<usize> {
    if let Some(ref mentions) = conf.on_mention {
        for mention in mentions {
            if !content.starts_with(&mention[..]) {
                continue;
            }

            return Some(mention.len());
        }
    }

    None
}
<commit_msg>Check if message starts with dynamic_prefix result<commit_after>use std::sync::Arc;
use super::Configuration;
use ::client::Context;
use ::model::Message;
use ::model::Permissions;
use std::collections::HashMap;

pub type Check = Fn(&mut Context, &Message) -> bool + Send + Sync + 'static;
pub type Exec = Fn(&mut Context, &Message, Vec<String>) -> Result<(), String> + Send + Sync + 'static;
pub type Help = Fn(&mut Context, &Message, HashMap<String, Arc<CommandGroup>>, &[String]) -> Result<(), String> + Send + Sync + 'static;
pub type BeforeHook = Fn(&mut Context, &Message, &String) -> bool + Send + Sync + 'static;
pub type AfterHook = Fn(&mut Context, &Message, &String, Result<(), String>) + Send + Sync + 'static;
#[doc(hidden)]
pub type InternalCommand = Arc<Command>;
pub type PrefixCheck = Fn(&mut Context) -> Option<String> + Send + Sync + 'static;

#[doc(hidden)]
pub enum CommandOrAlias {
    Alias(String),
    Command(InternalCommand),
}

/// Command function type. Allows to access internal framework things inside
/// your commands.
pub enum CommandType {
    StringResponse(String),
    Basic(Box<Exec>),
    WithCommands(Box<Help>),
}

#[derive(Default)]
pub struct CommandGroup {
    pub prefix: Option<String>,
    pub commands: HashMap<String, CommandOrAlias>,
}

/// Command struct used to store commands internally.
pub struct Command {
    /// A set of checks to be called prior to executing the command. The checks
    /// will short-circuit on the first check that returns `false`.
    pub checks: Vec<Box<Check>>,
    /// Function called when the command is called.
    pub exec: CommandType,
    /// Ratelimit bucket.
    pub bucket: Option<String>,
    /// Command description, used by other commands.
    pub desc: Option<String>,
    /// Example arguments, used by other commands.
    pub example: Option<String>,
    /// Command usage schema, used by other commands.
    pub usage: Option<String>,
    /// Whether arguments should be parsed using quote parser or not.
    pub use_quotes: bool,
    /// Minumum amount of arguments that should be passed.
    pub min_args: Option<i32>,
    /// Maximum amount of arguments that can be passed.
    pub max_args: Option<i32>,
    /// Permissions required to use this command.
    pub required_permissions: Permissions,
    /// Whether command should be displayed in help list or not, used by other commands.
    pub help_available: bool,
    /// Whether command can be used only privately or not.
    pub dm_only: bool,
    /// Whether command can be used only in guilds or not.
    pub guild_only: bool,
    /// Whether command can only be used by owners or not.
    pub owners_only: bool,
    #[doc(hidden)]
    pub aliases: Vec<String>,
}

impl Command {
    pub fn new<F>(f: F) -> Self
        where F: Fn(&mut Context, &Message, Vec<String>) -> Result<(), String> + Send + Sync + 'static {
        Command {
            aliases: Vec::new(),
            checks: Vec::default(),
            exec: CommandType::Basic(Box::new(f)),
            desc: None,
            usage: None,
            example: None,
            use_quotes: false,
            dm_only: false,
            bucket: None,
            guild_only: false,
            help_available: true,
            min_args: None,
            max_args: None,
            owners_only: false,
            required_permissions: Permissions::empty(),
        }
    }
}

pub fn positions(ctx: &mut Context, content: &str, conf: &Configuration) -> Option<Vec<usize>> {
    if !conf.prefixes.is_empty() || conf.dynamic_prefix.is_some() {
        // Find out if they were mentioned. If not, determine if the prefix
        // was used. If not, return None.
        let mut positions: Vec<usize> = vec![];

        if let Some(mention_end) = find_mention_end(content, conf) {
            positions.push(mention_end);
        } else if let Some(ref func) = conf.dynamic_prefix {
            if let Some(x) = func(ctx) {
                if content.starts_with(&x) {
                    positions.push(x.len());
                }
            } else {
                for n in conf.prefixes.clone() {
                    if content.starts_with(&n) {
                        positions.push(n.len());
                    }
                }
            }
        } else {
            for n in conf.prefixes.clone() {
                if content.starts_with(&n) {
                    positions.push(n.len());
                }
            }
        };

        if positions.is_empty() {
            return None;
        }

        if conf.allow_whitespace {
            let pos = *unsafe { positions.get_unchecked(0) };

            positions.insert(0, pos + 1);
        }

        Some(positions)
    } else if conf.on_mention.is_some() {
        match find_mention_end(content, conf) {
            Some(mention_end) => {
                let mut positions = vec![mention_end];

                if conf.allow_whitespace {
                    positions.insert(0, mention_end + 1);
                }

                Some(positions)
            },
            None => None,
        }
    } else {
        None
    }
}

fn find_mention_end(content: &str, conf: &Configuration) -> Option<usize> {
    if let Some(ref mentions) = conf.on_mention {
        for mention in mentions {
            if !content.starts_with(&mention[..]) {
                continue;
            }

            return Some(mention.len());
        }
    }

    None
}
<|endoftext|>


<commit_before>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use rustc::session::Session;
use rustc::util::common::{ProfQDumpParams, ProfileQueriesMsg, profq_msg, profq_set_chan};
use std::sync::mpsc::{Receiver};
use std::io::{Write};
use rustc::dep_graph::{DepNode};
use std::time::{Duration, Instant};

pub mod trace;

/// begin a profile thread, if not already running
pub fn begin(sess: &Session) {
    use std::thread;
    use std::sync::mpsc::{channel};
    let (tx, rx) = channel();
    if profq_set_chan(sess, tx) {
        thread::spawn(move || profile_queries_thread(rx));
    }
}

/// dump files with profiling information to the given base path, and
/// wait for this dump to complete.
///
/// wraps the RPC (send/recv channel logic) of requesting a dump.
pub fn dump(sess: &Session, path: String) {
    use std::sync::mpsc::{channel};
    let (tx, rx) = channel();
    let params = ProfQDumpParams {
        path,
        ack: tx,
        // FIXME: Add another compiler flag to toggle whether this log
        // is written; false for now
        dump_profq_msg_log: true,
    };
    profq_msg(sess, ProfileQueriesMsg::Dump(params));
    let _ = rx.recv().unwrap();
}

// State for parsing recursive trace structure in separate thread, via messages
#[derive(Clone, Eq, PartialEq)]
enum ParseState {
    // No (local) parse state; may be parsing a tree, focused on a
    // sub-tree that could be anything.
    Clear,
    // Have Query information from the last message
    HaveQuery(trace::Query, Instant),
    // Have "time-begin" information from the last message (doit flag, and message)
    HaveTimeBegin(String, Instant),
    // Have "task-begin" information from the last message
    HaveTaskBegin(DepNode, Instant),
}
struct StackFrame {
    pub parse_st: ParseState,
    pub traces:   Vec<trace::Rec>,
}

fn total_duration(traces: &[trace::Rec]) -> Duration {
    let mut sum : Duration = Duration::new(0, 0);
    for t in traces.iter() { sum += t.dur_total; }
    return sum
}

// profiling thread; retains state (in local variables) and dump traces, upon request.
fn profile_queries_thread(r: Receiver<ProfileQueriesMsg>) {
    use self::trace::*;
    use std::fs::File;
    use std::time::{Instant};

    let mut profq_msgs: Vec<ProfileQueriesMsg> = vec![];
    let mut frame: StackFrame = StackFrame { parse_st: ParseState::Clear, traces: vec![] };
    let mut stack: Vec<StackFrame> = vec![];
    loop {
        let msg = r.recv();
        if let Err(_recv_err) = msg {
            // FIXME: Perhaps do something smarter than simply quitting?
            break
        };
        let msg = msg.unwrap();
        debug!("profile_queries_thread: {:?}", msg);

        // Meta-level versus _actual_ queries messages
        match msg {
            ProfileQueriesMsg::Halt => return,
            ProfileQueriesMsg::Dump(params) => {
                assert!(stack.is_empty());
                assert!(frame.parse_st == ParseState::Clear);
                {
                    // write log of all messages
                    if params.dump_profq_msg_log {
                        let mut log_file =
                            File::create(format!("{}.log.txt", params.path)).unwrap();
                        for m in profq_msgs.iter() {
                            writeln!(&mut log_file, "{:?}", m).unwrap()
                        };
                    }

                    // write HTML file, and counts file
                    let html_path = format!("{}.html", params.path);
                    let mut html_file = File::create(&html_path).unwrap();

                    let counts_path = format!("{}.counts.txt", params.path);
                    let mut counts_file = File::create(&counts_path).unwrap();

                    writeln!(html_file,
                        "\n\n<link rel=\"stylesheet\" type=\"text/css\" href=\"{}\">",
                        "profile_queries.css").unwrap();
                    writeln!(html_file, "\n\n").unwrap();
                    trace::write_traces(&mut html_file, &mut counts_file, &frame.traces);
                    writeln!(html_file, "\n").unwrap();

                    let ack_path = format!("{}.ack", params.path);
                    let ack_file = File::create(&ack_path).unwrap();
                    drop(ack_file);

                    // Tell main thread that we are done, e.g., so it can exit
                    params.ack.send(()).unwrap();
                }
                continue
            }
            // Actual query message:
            msg => {
                // Record msg in our log
                profq_msgs.push(msg.clone());
                // Respond to the message, knowing that we've already handled Halt and Dump, above.
                match (frame.parse_st.clone(), msg) {
                    (_,ProfileQueriesMsg::Halt) => unreachable!(),
                    (_,ProfileQueriesMsg::Dump(_)) => unreachable!(),

                    // Parse State: Clear
                    (ParseState::Clear,
                     ProfileQueriesMsg::QueryBegin(span, querymsg)) => {
                        let start = Instant::now();
                        frame.parse_st = ParseState::HaveQuery
                            (Query { span, msg: querymsg }, start)
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::CacheHit) => {
                        panic!("parse error: unexpected CacheHit; expected QueryBegin")
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::ProviderBegin) => {
                        panic!("parse error: expected QueryBegin before beginning a provider")
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::ProviderEnd) => {
                        let provider_extent = frame.traces;
                        match stack.pop() {
                            None =>
                                panic!("parse error: expected a stack frame; found an empty stack"),
                            Some(old_frame) => {
                                match old_frame.parse_st {
                                    ParseState::HaveQuery(q, start) => {
                                        let duration = start.elapsed();
                                        frame = StackFrame{
                                            parse_st:ParseState::Clear,
                                            traces:old_frame.traces
                                        };
                                        let dur_extent = total_duration(&provider_extent);
                                        let trace = Rec {
                                            effect: Effect::QueryBegin(q, CacheCase::Miss),
                                            extent: Box::new(provider_extent),
                                            start: start,
                                            dur_self: duration - dur_extent,
                                            dur_total: duration,
                                        };
                                        frame.traces.push( trace );
                                    },
                                    _ => panic!("internal parse error: malformed parse stack")
                                }
                            }
                        }
                    },

                    (ParseState::Clear,
                     ProfileQueriesMsg::TimeBegin(msg)) => {
                        let start = Instant::now();
                        frame.parse_st = ParseState::HaveTimeBegin(msg, start);
                        stack.push(frame);
                        frame = StackFrame{parse_st:ParseState::Clear, traces:vec![]};
                    },
                    (_, ProfileQueriesMsg::TimeBegin(_)) =>
                        panic!("parse error; did not expect time begin here"),

                    (ParseState::Clear,
                     ProfileQueriesMsg::TimeEnd) => {
                        let provider_extent = frame.traces;
                        match stack.pop() {
                            None =>
                                panic!("parse error: expected a stack frame; found an empty stack"),
                            Some(old_frame) => {
                                match old_frame.parse_st {
                                    ParseState::HaveTimeBegin(msg, start) => {
                                        let duration = start.elapsed();
                                        frame = StackFrame{
                                            parse_st:ParseState::Clear,
                                            traces:old_frame.traces
                                        };
                                        let dur_extent = total_duration(&provider_extent);
                                        let trace = Rec {
                                            effect: Effect::TimeBegin(msg),
                                            extent: Box::new(provider_extent),
                                            start: start,
                                            dur_total: duration,
                                            dur_self: duration - dur_extent,
                                        };
                                        frame.traces.push( trace );
                                    },
                                    _ => panic!("internal parse error: malformed parse stack")
                                }
                            }
                        }
                    },
                    (_, ProfileQueriesMsg::TimeEnd) => { panic!("parse error") }

                    (ParseState::Clear,
                     ProfileQueriesMsg::TaskBegin(key)) => {
                        let start = Instant::now();
                        frame.parse_st = ParseState::HaveTaskBegin(key, start);
                        stack.push(frame);
                        frame = StackFrame{parse_st:ParseState::Clear, traces:vec![]};
                    },
                    (_, ProfileQueriesMsg::TaskBegin(_)) =>
                        panic!("parse error; did not expect time begin here"),

                    (ParseState::Clear,
                     ProfileQueriesMsg::TaskEnd) => {
                        let provider_extent = frame.traces;
                        match stack.pop() {
                            None =>
                                panic!("parse error: expected a stack frame; found an empty stack"),
                            Some(old_frame) => {
                                match old_frame.parse_st {
                                    ParseState::HaveTaskBegin(key, start) => {
                                        let duration = start.elapsed();
                                        frame = StackFrame{
                                            parse_st:ParseState::Clear,
                                            traces:old_frame.traces
                                        };
                                        let dur_extent = total_duration(&provider_extent);
                                        let trace = Rec {
                                            effect: Effect::TaskBegin(key),
                                            extent: Box::new(provider_extent),
                                            start: start,
                                            dur_total: duration,
                                            dur_self: duration - dur_extent,
                                        };
                                        frame.traces.push( trace );
                                    },
                                    _ => panic!("internal parse error: malformed parse stack")
                                }
                            }
                        }
                    },
                    (_, ProfileQueriesMsg::TaskEnd) => { panic!("parse error") }

                    // Parse State: HaveQuery
                    (ParseState::HaveQuery(q,start),
                     ProfileQueriesMsg::CacheHit) => {
                        let duration = start.elapsed();
                        let trace : Rec = Rec{
                            effect: Effect::QueryBegin(q, CacheCase::Hit),
                            extent: Box::new(vec![]),
                            start: start,
                            dur_self: duration,
                            dur_total: duration,
                        };
                        frame.traces.push( trace );
                        frame.parse_st = ParseState::Clear;
                    },
                    (ParseState::HaveQuery(_,_),
                     ProfileQueriesMsg::ProviderBegin) => {
                        stack.push(frame);
                        frame = StackFrame{parse_st:ParseState::Clear, traces:vec![]};
                    },

                    // Parse errors:

                    (ParseState::HaveQuery(q,_),
                     ProfileQueriesMsg::ProviderEnd) => {
                        panic!("parse error: unexpected ProviderEnd; \
                                expected something else to follow BeginQuery for {:?}", q)
                    },
                    (ParseState::HaveQuery(q1,_),
                     ProfileQueriesMsg::QueryBegin(span2,querymsg2)) => {
                        panic!("parse error: unexpected QueryBegin; \
                                earlier query is unfinished: {:?} and now {:?}",
                               q1, Query{span:span2, msg:querymsg2})
                    },

                    (ParseState::HaveTimeBegin(_, _), _) => {
                        unreachable!()
                    },
                    (ParseState::HaveTaskBegin(_, _), _) => {
                        unreachable!()
                    },
                }
            }
        }
    }
}
<commit_msg>Rollup merge of #56918 - ljedrz:profiler_nits, r=wesleywiser<commit_after>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use rustc::session::Session;
use rustc::util::common::{ProfQDumpParams, ProfileQueriesMsg, profq_msg, profq_set_chan};
use std::sync::mpsc::{Receiver};
use std::io::{Write};
use rustc::dep_graph::{DepNode};
use std::time::{Duration, Instant};

pub mod trace;

/// begin a profile thread, if not already running
pub fn begin(sess: &Session) {
    use std::thread;
    use std::sync::mpsc::{channel};
    let (tx, rx) = channel();
    if profq_set_chan(sess, tx) {
        thread::spawn(move || profile_queries_thread(rx));
    }
}

/// dump files with profiling information to the given base path, and
/// wait for this dump to complete.
///
/// wraps the RPC (send/recv channel logic) of requesting a dump.
pub fn dump(sess: &Session, path: String) {
    use std::sync::mpsc::{channel};
    let (tx, rx) = channel();
    let params = ProfQDumpParams {
        path,
        ack: tx,
        // FIXME: Add another compiler flag to toggle whether this log
        // is written; false for now
        dump_profq_msg_log: true,
    };
    profq_msg(sess, ProfileQueriesMsg::Dump(params));
    let _ = rx.recv().unwrap();
}

// State for parsing recursive trace structure in separate thread, via messages
#[derive(Clone, Eq, PartialEq)]
enum ParseState {
    // No (local) parse state; may be parsing a tree, focused on a
    // sub-tree that could be anything.
    Clear,
    // Have Query information from the last message
    HaveQuery(trace::Query, Instant),
    // Have "time-begin" information from the last message (doit flag, and message)
    HaveTimeBegin(String, Instant),
    // Have "task-begin" information from the last message
    HaveTaskBegin(DepNode, Instant),
}
struct StackFrame {
    pub parse_st: ParseState,
    pub traces:   Vec<trace::Rec>,
}

fn total_duration(traces: &[trace::Rec]) -> Duration {
    Duration::new(0, 0) + traces.iter().map(|t| t.dur_total).sum()
}

// profiling thread; retains state (in local variables) and dump traces, upon request.
fn profile_queries_thread(r: Receiver<ProfileQueriesMsg>) {
    use self::trace::*;
    use std::fs::File;
    use std::time::{Instant};

    let mut profq_msgs: Vec<ProfileQueriesMsg> = vec![];
    let mut frame: StackFrame = StackFrame { parse_st: ParseState::Clear, traces: vec![] };
    let mut stack: Vec<StackFrame> = vec![];
    loop {
        let msg = r.recv();
        if let Err(_recv_err) = msg {
            // FIXME: Perhaps do something smarter than simply quitting?
            break
        };
        let msg = msg.unwrap();
        debug!("profile_queries_thread: {:?}", msg);

        // Meta-level versus _actual_ queries messages
        match msg {
            ProfileQueriesMsg::Halt => return,
            ProfileQueriesMsg::Dump(params) => {
                assert!(stack.is_empty());
                assert!(frame.parse_st == ParseState::Clear);

                // write log of all messages
                if params.dump_profq_msg_log {
                    let mut log_file =
                        File::create(format!("{}.log.txt", params.path)).unwrap();
                    for m in profq_msgs.iter() {
                        writeln!(&mut log_file, "{:?}", m).unwrap()
                    };
                }

                // write HTML file, and counts file
                let html_path = format!("{}.html", params.path);
                let mut html_file = File::create(&html_path).unwrap();

                let counts_path = format!("{}.counts.txt", params.path);
                let mut counts_file = File::create(&counts_path).unwrap();

                writeln!(html_file,
                    "\n\n<link rel=\"stylesheet\" type=\"text/css\" href=\"{}\">",
                    "profile_queries.css").unwrap();
                writeln!(html_file, "\n\n").unwrap();
                trace::write_traces(&mut html_file, &mut counts_file, &frame.traces);
                writeln!(html_file, "\n").unwrap();

                let ack_path = format!("{}.ack", params.path);
                let ack_file = File::create(&ack_path).unwrap();
                drop(ack_file);

                // Tell main thread that we are done, e.g., so it can exit
                params.ack.send(()).unwrap();
            }
            // Actual query message:
            msg => {
                // Record msg in our log
                profq_msgs.push(msg.clone());
                // Respond to the message, knowing that we've already handled Halt and Dump, above.
                match (frame.parse_st.clone(), msg) {
                    (_, ProfileQueriesMsg::Halt) | (_, ProfileQueriesMsg::Dump(_)) => {
                        unreachable!();
                    },
                    // Parse State: Clear
                    (ParseState::Clear,
                     ProfileQueriesMsg::QueryBegin(span, querymsg)) => {
                        let start = Instant::now();
                        frame.parse_st = ParseState::HaveQuery
                            (Query { span, msg: querymsg }, start)
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::CacheHit) => {
                        panic!("parse error: unexpected CacheHit; expected QueryBegin")
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::ProviderBegin) => {
                        panic!("parse error: expected QueryBegin before beginning a provider")
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::ProviderEnd) => {
                        let provider_extent = frame.traces;
                        match stack.pop() {
                            None =>
                                panic!("parse error: expected a stack frame; found an empty stack"),
                            Some(old_frame) => {
                                match old_frame.parse_st {
                                    ParseState::HaveQuery(q, start) => {
                                        let duration = start.elapsed();
                                        frame = StackFrame{
                                            parse_st: ParseState::Clear,
                                            traces: old_frame.traces
                                        };
                                        let dur_extent = total_duration(&provider_extent);
                                        let trace = Rec {
                                            effect: Effect::QueryBegin(q, CacheCase::Miss),
                                            extent: Box::new(provider_extent),
                                            start: start,
                                            dur_self: duration - dur_extent,
                                            dur_total: duration,
                                        };
                                        frame.traces.push( trace );
                                    },
                                    _ => panic!("internal parse error: malformed parse stack")
                                }
                            }
                        }
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::TimeBegin(msg)) => {
                        let start = Instant::now();
                        frame.parse_st = ParseState::HaveTimeBegin(msg, start);
                        stack.push(frame);
                        frame = StackFrame{parse_st: ParseState::Clear, traces: vec![]};
                    },
                    (_, ProfileQueriesMsg::TimeBegin(_)) => {
                        panic!("parse error; did not expect time begin here");
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::TimeEnd) => {
                        let provider_extent = frame.traces;
                        match stack.pop() {
                            None =>
                                panic!("parse error: expected a stack frame; found an empty stack"),
                            Some(old_frame) => {
                                match old_frame.parse_st {
                                    ParseState::HaveTimeBegin(msg, start) => {
                                        let duration = start.elapsed();
                                        frame = StackFrame{
                                            parse_st: ParseState::Clear,
                                            traces: old_frame.traces
                                        };
                                        let dur_extent = total_duration(&provider_extent);
                                        let trace = Rec {
                                            effect: Effect::TimeBegin(msg),
                                            extent: Box::new(provider_extent),
                                            start: start,
                                            dur_total: duration,
                                            dur_self: duration - dur_extent,
                                        };
                                        frame.traces.push( trace );
                                    },
                                    _ => panic!("internal parse error: malformed parse stack")
                                }
                            }
                        }
                    },
                    (_, ProfileQueriesMsg::TimeEnd) => {
                        panic!("parse error")
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::TaskBegin(key)) => {
                        let start = Instant::now();
                        frame.parse_st = ParseState::HaveTaskBegin(key, start);
                        stack.push(frame);
                        frame = StackFrame{ parse_st: ParseState::Clear, traces: vec![] };
                    },
                    (_, ProfileQueriesMsg::TaskBegin(_)) => {
                        panic!("parse error; did not expect time begin here");
                    },
                    (ParseState::Clear,
                     ProfileQueriesMsg::TaskEnd) => {
                        let provider_extent = frame.traces;
                        match stack.pop() {
                            None =>
                                panic!("parse error: expected a stack frame; found an empty stack"),
                            Some(old_frame) => {
                                match old_frame.parse_st {
                                    ParseState::HaveTaskBegin(key, start) => {
                                        let duration = start.elapsed();
                                        frame = StackFrame{
                                            parse_st: ParseState::Clear,
                                            traces: old_frame.traces
                                        };
                                        let dur_extent = total_duration(&provider_extent);
                                        let trace = Rec {
                                            effect: Effect::TaskBegin(key),
                                            extent: Box::new(provider_extent),
                                            start: start,
                                            dur_total: duration,
                                            dur_self: duration - dur_extent,
                                        };
                                        frame.traces.push( trace );
                                    },
                                    _ => panic!("internal parse error: malformed parse stack")
                                }
                            }
                        }
                    },
                    (_, ProfileQueriesMsg::TaskEnd) => {
                        panic!("parse error")
                    },
                    // Parse State: HaveQuery
                    (ParseState::HaveQuery(q,start),
                     ProfileQueriesMsg::CacheHit) => {
                        let duration = start.elapsed();
                        let trace : Rec = Rec{
                            effect: Effect::QueryBegin(q, CacheCase::Hit),
                            extent: Box::new(vec![]),
                            start: start,
                            dur_self: duration,
                            dur_total: duration,
                        };
                        frame.traces.push( trace );
                        frame.parse_st = ParseState::Clear;
                    },
                    (ParseState::HaveQuery(_, _),
                     ProfileQueriesMsg::ProviderBegin) => {
                        stack.push(frame);
                        frame = StackFrame{ parse_st: ParseState::Clear, traces: vec![] };
                    },

                    // Parse errors:

                    (ParseState::HaveQuery(q, _),
                     ProfileQueriesMsg::ProviderEnd) => {
                        panic!("parse error: unexpected ProviderEnd; \
                                expected something else to follow BeginQuery for {:?}", q)
                    },
                    (ParseState::HaveQuery(q1, _),
                     ProfileQueriesMsg::QueryBegin(span2, querymsg2)) => {
                        panic!("parse error: unexpected QueryBegin; \
                                earlier query is unfinished: {:?} and now {:?}",
                               q1, Query{span:span2, msg: querymsg2})
                    },
                    (ParseState::HaveTimeBegin(_, _), _) => {
                        unreachable!()
                    },
                    (ParseState::HaveTaskBegin(_, _), _) => {
                        unreachable!()
                    },
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::iter;
use std::fmt::{self, Debug};
use std::cmp::{Eq, PartialEq};
use std::ops::{Deref, DerefMut};
use rustwlc::{Geometry, Size, Point};
use rustwlc::render::{write_pixels, wlc_pixel_format};
use cairo::{self, Context, ImageSurface, Format, Operator, Status, SolidPattern};
use cairo::prelude::{SurfaceExt};

use super::draw::BaseDraw;

/// The borders of a container.
///
/// This type just deals with rendering,
#[derive(Clone)]
pub struct Borders {
    /// The surface that contains the bytes we give to wlc to draw.
    surface: ImageSurface,
    /// The geometry where the buffer is written.
    ///
    /// Should correspond with the geometry of the container.
    pub geometry: Geometry
}

impl Borders {
    pub fn new(geometry: Geometry) -> Self {
        let Size { w, h } = geometry.size;
        let stride = calculate_stride(w) as i32;
        let data: Vec<u8> = iter::repeat(100).take(w as usize * stride as usize).collect();
        let buffer = data.into_boxed_slice();
        let surface = ImageSurface::create_for_data(buffer,
                                                    drop_data,
                                                    Format::ARgb32,
                                                    h as i32,
                                                    w as i32,
                                                    stride);
        Borders {
            surface: surface,
            geometry: geometry
        }
    }

    pub fn render(&mut self) {
        let mut buffer = self.surface.get_data()
            .expect("Could not get border surface buffer");
        write_pixels(wlc_pixel_format::WLC_RGBA8888, self.geometry, &mut buffer);
    }

    /// Updates the position at which we render the buffer.
    ///
    /// Since this is a common operation, and does not resize the buffer,
    /// this is a method is here instead of in a draw struct
    /// for ergonomic and performance reasons.
    ///
    /// # Rendering
    /// Note that this not actually cause a rerender, but sets the internal
    /// flag so that wlc can render it on the next render cycle.
    pub fn update_pos(&mut self, origin: Point) {
        self.geometry.origin = origin;
    }

    /// Enables Cairo drawing capabilities on the borders.
    ///
    /// While drawing with Cairo, the underlying surface cannot be read.
    /// Borders is consumed in order for this to be checked at compile time.
    pub fn enable_cairo(self) -> Result<BaseDraw, cairo::Status> {
        let cairo = Context::new(&self.surface);
        match cairo.status() {
            cairo::Status::Success => {},
            err => return Err(err)
        }
        cairo.set_operator(Operator::Source);
        match cairo.status() {
            cairo::Status::Success => {},
            err => return Err(err)
        }
        Ok(BaseDraw::new(self, cairo))
    }
}

impl Debug for Borders {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Debug")
            .field("geometry", &self.geometry as &Debug)
            .finish()
    }
}

impl PartialEq for Borders {
    fn eq(&self, other: &Borders) -> bool {
        self.geometry == other.geometry
    }
}

impl Eq for Borders {}

unsafe impl Send for Borders {}
unsafe impl Sync for Borders {}

/// Calculates the stride for ARgb32 encoded buffers
fn calculate_stride(width: u32) -> u32 {
    // function stolen from CAIRO_STRIDE_FOR_WIDTH macro in cairoint.h
    // Can be found in the most recent version of the cairo source
    (32 * width + 7) / 8 + 4 & !4
}

#[allow(dead_code)]
fn drop_data(_: Box<[u8]>) {}
<commit_msg>Fixed buffer creation for borders<commit_after>use std::iter;
use std::fmt::{self, Debug};
use std::cmp::{Eq, PartialEq};
use std::ops::{Deref, DerefMut};
use rustwlc::{Geometry, Size, Point};
use rustwlc::render::{write_pixels, wlc_pixel_format, BITS_PER_PIXEL};
use cairo::{self, Context, ImageSurface, Format, Operator, Status, SolidPattern};
use cairo::prelude::{SurfaceExt};

use super::draw::BaseDraw;

/// The borders of a container.
///
/// This type just deals with rendering,
#[derive(Clone)]
pub struct Borders {
    /// The surface that contains the bytes we give to wlc to draw.
    surface: ImageSurface,
    /// The geometry where the buffer is written.
    ///
    /// Should correspond with the geometry of the container.
    pub geometry: Geometry
}

impl Borders {
    pub fn new(geometry: Geometry) -> Self {
        let Size { w, h } = geometry.size;
        let stride = calculate_stride(w) as i32;
        let data: Vec<u8> = iter::repeat(100).take(h as usize * stride as usize).collect();
        let buffer = data.into_boxed_slice();
        let surface = ImageSurface::create_for_data(buffer,
                                                    drop_data,
                                                    Format::ARgb32,
                                                    w as i32,
                                                    h as i32,
                                                    stride);
        Borders {
            surface: surface,
            geometry: geometry
        }
    }

    pub fn render(&mut self) {
        let mut buffer = self.surface.get_data()
            .expect("Could not get border surface buffer");
        write_pixels(wlc_pixel_format::WLC_RGBA8888, self.geometry, &mut buffer);
    }

    /// Updates the position at which we render the buffer.
    ///
    /// Since this is a common operation, and does not resize the buffer,
    /// this is a method is here instead of in a draw struct
    /// for ergonomic and performance reasons.
    ///
    /// # Rendering
    /// Note that this not actually cause a rerender, but sets the internal
    /// flag so that wlc can render it on the next render cycle.
    pub fn update_pos(&mut self, origin: Point) {
        self.geometry.origin = origin;
    }

    /// Enables Cairo drawing capabilities on the borders.
    ///
    /// While drawing with Cairo, the underlying surface cannot be read.
    /// Borders is consumed in order for this to be checked at compile time.
    pub fn enable_cairo(self) -> Result<BaseDraw, cairo::Status> {
        let cairo = Context::new(&self.surface);
        match cairo.status() {
            cairo::Status::Success => {},
            err => return Err(err)
        }
        cairo.set_operator(Operator::Source);
        match cairo.status() {
            cairo::Status::Success => {},
            err => return Err(err)
        }
        Ok(BaseDraw::new(self, cairo))
    }
}

impl Debug for Borders {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Debug")
            .field("geometry", &self.geometry as &Debug)
            .finish()
    }
}

impl PartialEq for Borders {
    fn eq(&self, other: &Borders) -> bool {
        self.geometry == other.geometry
    }
}

impl Eq for Borders {}

unsafe impl Send for Borders {}
unsafe impl Sync for Borders {}

/// Calculates the stride for ARgb32 encoded buffers
fn calculate_stride(width: u32) -> u32 {
    // function stolen from CAIRO_STRIDE_FOR_WIDTH macro in cairoint.h
    // Can be found in the most recent version of the cairo source
    let stride_alignment = ::std::mem::size_of::<u32>() as u32;
    ((BITS_PER_PIXEL * width + 7 ) / 8 + (stride_alignment - 1))  & (stride_alignment.overflowing_neg().0)
}

#[allow(dead_code)]
fn drop_data(_: Box<[u8]>) {}
<|endoftext|>