<commit_before>extern crate clap;
extern crate env_logger;
#[macro_use]
extern crate log;
#[cfg(test)]
extern crate tempdir;
extern crate tuf as _tuf;
extern crate url;

use clap::{App, AppSettings, SubCommand, Arg, ArgMatches};
use std::path::PathBuf;
use _tuf::{Tuf, Config, Error};
use url::Url;

// TODO logging

fn main() {
    let matches = parser().get_matches();
    env_logger::init().unwrap();

    match run_main(matches) {
        Ok(()) => std::process::exit(0),
        Err(e) => {
            error!("{:?}", e);
            std::process::exit(1);
        }
    }
}

fn run_main(matches: ArgMatches) -> Result<(), Error> {
    let config = Config::build().url(Url::parse(matches.value_of("url").unwrap())?)
        .local_path(PathBuf::from(matches.value_of("path").unwrap()))
        .finish()?;

    if let Some(_) = matches.subcommand_matches("init") {
        let path = PathBuf::from(matches.value_of("path").unwrap());
        cmd_init(&path)
    } else if let Some(_) = matches.subcommand_matches("list") {
        let mut tuf = Tuf::new(config)?;
        cmd_list(&mut tuf)
    } else if let Some(_) = matches.subcommand_matches("update") {
        let mut tuf = Tuf::new(config)?;
        cmd_update(&mut tuf)
    } else if let Some(matches) = matches.subcommand_matches("verify") {
        let mut tuf = Tuf::new(config)?;
        cmd_verify(&mut tuf, matches.value_of("target").unwrap())
    } else {
        unreachable!() // because of AppSettings::SubcommandRequiredElseHelp
    }
}

fn url_validator(url: String) -> Result<(), String> {
    Url::parse(&url)
        .map(|_| ())
        .map_err(|_| "URL was not valid".into())
}

fn parser<'a, 'b>() -> App<'a, 'b> {
    App::new("tuf")
        .version(env!("CARGO_PKG_VERSION"))
        .about("CLI tool for verifying TUF metadata and downloading targets")
        .settings(&[AppSettings::SubcommandRequiredElseHelp])
        .arg(Arg::with_name("verbose")
            .short("v")
            .long("verbose")
            .multiple(true)
            .help("Increase the verbosity of output to stderr"))
        .arg(Arg::with_name("url")
            .short("U")
            .long("url")
            .takes_value(true)
            .required(true)
            .validator(url_validator)
            .help("URL of the TUF repo (local or remote)"))
        .arg(Arg::with_name("path")
            .short("p")
            .long("path")
            .takes_value(true)
            .required(true)
            .help("Local path the TUF repo"))
        .subcommand(SubCommand::with_name("init").about("Initializes a new TUF repo"))
        .subcommand(SubCommand::with_name("list").about("Lists available targets"))
        .subcommand(SubCommand::with_name("update").about("Updates metadata from remotes"))
        .subcommand(SubCommand::with_name("verify")
            .about("Verifies a target")
            .arg(Arg::with_name("target")
                .takes_value(true)
                .required(true)
                .help("The full (non-local) path of the target to verify")))
}

fn cmd_init(local_path: &PathBuf) -> Result<(), Error> {
    Tuf::initialize(local_path)
}

fn cmd_list(tuf: &mut Tuf) -> Result<(), Error> {
    let mut targets = tuf.list_targets();
    targets.sort();

    for target in targets.iter() {
        println!("{}", target);
    }

    Ok(())
}

fn cmd_update(tuf: &mut Tuf) -> Result<(), Error> {
    tuf.update()
}

fn cmd_verify(tuf: &mut Tuf, target: &str) -> Result<(), Error> {
    tuf.fetch_target(target)
        .map(|_| ())
}

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::{self, DirBuilder};
    use std::path::{Path, PathBuf};
    use tempdir::TempDir;
    use _tuf::util;

    fn vector_path() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("tests")
            .join("tuf-test-vectors")
            .join("vectors")
            .join("001")
            .join("repo")
    }

    #[test]
    fn test_clap() {
        let _ = parser();
    }

    fn init_temp(temp: &Path) {
        let dir = PathBuf::from("metadata").join("current");
        DirBuilder::new()
            .recursive(true)
            .create(temp.join(dir.clone()))
            .expect(&format!("couldn't create path {}:", temp.join(dir).to_string_lossy()));

        let copy_path = vector_path().join("root.json");
        fs::copy(copy_path,
                 temp.join("metadata").join("current").join("root.json"))
            .expect(&format!("copy failed for target"));
    }

    #[test]
    fn run_it() {
        let temp = TempDir::new("rust-tuf").expect("couldn't make temp dir");
        init_temp(temp.path());

        let matches = parser()
            .get_matches_from_safe(vec!["tuf",
                                        "--url",
                                        &util::path_to_url(&vector_path())
                                            .expect("bad path")
                                            .to_string(),
                                        "--path",
                                        temp.path().to_str().expect("path not utf-8"),
                                        "init"])
            .expect("parse error");
        assert_eq!(run_main(matches), Ok(()));

        let matches = parser()
            .get_matches_from_safe(vec!["tuf",
                                        "--url",
                                        &util::path_to_url(&vector_path())
                                            .expect("bad path")
                                            .to_string(),
                                        "--path",
                                        temp.path().to_str().expect("path not utf-8"),
                                        "update"])
            .expect("parse error");
        assert_eq!(run_main(matches), Ok(()));

        let matches = parser()
            .get_matches_from_safe(vec!["tuf",
                                        "--url",
                                        &util::path_to_url(&vector_path())
                                            .expect("bad path")
                                            .to_string(),
                                        "--path",
                                        temp.path().to_str().expect("path not utf-8"),
                                        "verify",
                                        "targets/file.txt"])
            .expect("parse error");
        assert_eq!(run_main(matches), Ok(()));
    }
}
<commit_msg>fixed cli build<commit_after>extern crate clap;
extern crate env_logger;
#[macro_use]
extern crate log;
#[cfg(test)]
extern crate tempdir;
extern crate tuf as _tuf;
extern crate url;

use clap::{App, AppSettings, SubCommand, Arg, ArgMatches};
use std::path::PathBuf;
use _tuf::{Tuf, Config, Error};
use url::Url;

// TODO logging

fn main() {
    let matches = parser().get_matches();
    env_logger::init().unwrap();

    match run_main(matches) {
        Ok(()) => std::process::exit(0),
        Err(e) => {
            error!("{:?}", e);
            std::process::exit(1);
        }
    }
}

fn run_main(matches: ArgMatches) -> Result<(), Error> {
    let config = Config::build().url(Url::parse(matches.value_of("url").unwrap())?)
        .local_path(PathBuf::from(matches.value_of("path").unwrap()))
        .finish()?;

    if let Some(_) = matches.subcommand_matches("init") {
        let path = PathBuf::from(matches.value_of("path").unwrap());
        cmd_init(&path)
    } else if let Some(_) = matches.subcommand_matches("list") {
        let mut tuf = Tuf::new(config)?;
        cmd_list(&mut tuf)
    } else if let Some(_) = matches.subcommand_matches("update") {
        let mut tuf = Tuf::new(config)?;
        cmd_update(&mut tuf)
    } else if let Some(matches) = matches.subcommand_matches("verify") {
        let mut tuf = Tuf::new(config)?;
        cmd_verify(&mut tuf, matches.value_of("target").unwrap())
    } else {
        unreachable!() // because of AppSettings::SubcommandRequiredElseHelp
    }
}

fn url_validator(url: String) -> Result<(), String> {
    Url::parse(&url)
        .map(|_| ())
        .map_err(|_| "URL was not valid".into())
}

fn parser<'a, 'b>() -> App<'a, 'b> {
    App::new("tuf")
        .version(env!("CARGO_PKG_VERSION"))
        .about("CLI tool for verifying TUF metadata and downloading targets")
        .settings(&[AppSettings::SubcommandRequiredElseHelp])
        .arg(Arg::with_name("verbose")
            .short("v")
            .long("verbose")
            .multiple(true)
            .help("Increase the verbosity of output to stderr"))
        .arg(Arg::with_name("url")
            .short("U")
            .long("url")
            .takes_value(true)
            .required(true)
            .validator(url_validator)
            .help("URL of the TUF repo (local or remote)"))
        .arg(Arg::with_name("path")
            .short("p")
            .long("path")
            .takes_value(true)
            .required(true)
            .help("Local path the TUF repo"))
        .subcommand(SubCommand::with_name("init").about("Initializes a new TUF repo"))
        .subcommand(SubCommand::with_name("list").about("Lists available targets"))
        .subcommand(SubCommand::with_name("update").about("Updates metadata from remotes"))
        .subcommand(SubCommand::with_name("verify")
            .about("Verifies a target")
            .arg(Arg::with_name("target")
                .takes_value(true)
                .required(true)
                .help("The full (non-local) path of the target to verify")))
}

fn cmd_init(local_path: &PathBuf) -> Result<(), Error> {
    Tuf::initialize(local_path)
}

fn cmd_list(tuf: &mut Tuf) -> Result<(), Error> {
    let mut targets = tuf.list_targets();
    targets.sort();

    for target in targets.iter() {
        println!("{}", target);
    }

    Ok(())
}

fn cmd_update(tuf: &mut Tuf) -> Result<(), Error> {
    tuf.update()
}

fn cmd_verify(tuf: &mut Tuf, target: &str) -> Result<(), Error> {
    tuf.fetch_target(target)
        .map(|_| ())
}

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::{self, DirBuilder};
    use std::path::{Path, PathBuf};
    use tempdir::TempDir;
    use _tuf::util;

    fn vector_path() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("tests")
            .join("tuf-test-vectors")
            .join("tuf")
            .join("001")
            .join("repo")
    }

    #[test]
    fn test_clap() {
        let _ = parser();
    }

    fn init_temp(temp: &Path) {
        let dir = PathBuf::from("metadata").join("current");
        DirBuilder::new()
            .recursive(true)
            .create(temp.join(dir.clone()))
            .expect(&format!("couldn't create path {}:", temp.join(dir).to_string_lossy()));

        let copy_path = vector_path().join("root.json");
        fs::copy(copy_path,
                 temp.join("metadata").join("current").join("root.json"))
            .expect(&format!("copy failed for target"));
    }

    #[test]
    fn run_it() {
        let temp = TempDir::new("rust-tuf").expect("couldn't make temp dir");
        init_temp(temp.path());

        let matches = parser()
            .get_matches_from_safe(vec!["tuf",
                                        "--url",
                                        &util::path_to_url(&vector_path())
                                            .expect("bad path")
                                            .to_string(),
                                        "--path",
                                        temp.path().to_str().expect("path not utf-8"),
                                        "init"])
            .expect("parse error");
        assert_eq!(run_main(matches), Ok(()));

        let matches = parser()
            .get_matches_from_safe(vec!["tuf",
                                        "--url",
                                        &util::path_to_url(&vector_path())
                                            .expect("bad path")
                                            .to_string(),
                                        "--path",
                                        temp.path().to_str().expect("path not utf-8"),
                                        "update"])
            .expect("parse error");
        assert_eq!(run_main(matches), Ok(()));

        let matches = parser()
            .get_matches_from_safe(vec!["tuf",
                                        "--url",
                                        &util::path_to_url(&vector_path())
                                            .expect("bad path")
                                            .to_string(),
                                        "--path",
                                        temp.path().to_str().expect("path not utf-8"),
                                        "verify",
                                        "targets/file.txt"])
            .expect("parse error");
        assert_eq!(run_main(matches), Ok(()));
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 rust-fuzz developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

extern crate cargo_metadata;
extern crate docopt;
extern crate rustc_serialize;

use cargo_metadata::{metadata, Package};
use docopt::Docopt;
use std::{env, error, fs, io, path, process};
use std::io::Write;

const USAGE: &'static str = "
Cargo Fuzz

Usage:
  cargo fuzz --init
  cargo fuzz --fuzz-target TARGET
  cargo fuzz --add TARGET
  cargo fuzz --list
  cargo fuzz (-h | --help)

Options:
  -h --help              Show this screen.
  --init                 Initialize fuzz folder
  --fuzz-target TARGET   Run with given fuzz target in fuzz/fuzzers
  --add TARGET           Add a new fuzz target
  --list                 List the available fuzz targets
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_init: bool,
    flag_add: Option<String>,
    flag_fuzz_target: Option<String>,
    flag_list: bool,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let result = if args.flag_init {
        init_fuzz()
    } else if let Some(target) = args.flag_add {
        add_target(target)
    } else if let Some(target) = args.flag_fuzz_target {
        let result = run_target(target);
        if let Ok(success) = result {
            if success {
                // Can this ever happen?
                Ok(())
            } else {
                println!("Fuzzing found errors!");
                process::exit(-1)
            }
        } else {
            result.map(|_| ())
        }
    } else if args.flag_list {
        list_fuzz_targets()
            .map(|_| ())
    } else {
        println!("Invalid arguments. Usage:\n{}", USAGE);
        return;
    };
    if let Err(error) = result {
        writeln!(io::stderr(), "Error: {}", error)
            .expect("failed writing to stderr");
    }
}

fn list_fuzz_targets() -> Result<(), Box<error::Error>> {
    if !path::Path::new("./fuzz").is_dir() {
        return Err("Fuzzing crate has not been initialized. Run `cargo fuzz --init` to initialize it.".into());
    }
    env::set_current_dir("./fuzz")?;
    let package = get_package();
    for target in &package.targets {
        println!("{}", target.name);
    }
    Ok(())
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> Result<(), Box<error::Error>> {
    let me = get_package();

    fs::create_dir("./fuzz")?;
    fs::create_dir("./fuzz/fuzzers")?;

    let mut cargo = fs::File::create(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[package]
name = "{0}-fuzz"
version = "0.0.1"
authors = ["Automatically generated"]
publish = false

[dependencies.{0}]
path = ".."

# Prevent this from interfering with workspaces
[workspace]
members = ["."]

[[bin]]
name = "fuzzer_script_1"
path = "fuzzers/fuzzer_script_1.rs"
"#, me.name)?;

    let mut ignore = fs::File::create(path::Path::new("./fuzz/.gitignore"))?;

write!(ignore, r#"
target
libfuzzer
corpus
"#)?;

    let mut script = fs::File::create(path::Path::new("./fuzz/fuzzers/fuzzer_script_1.rs"))?;
    dummy_target(&mut script, &me)
}

/// Returns metadata for the Cargo package in the current directory
fn get_package() -> Package {
    // todo error handling
    let meta = metadata(None).unwrap();
    let mut p = env::current_dir().unwrap();
    p.push("Cargo.toml");
    let p = p.to_str().unwrap();
    meta.packages.into_iter().find(|package| package.manifest_path == p).unwrap()
}

/// If the package contains a library target, generate an `extern crate` line to link to it.
fn link_to_lib(pkg: &Package) -> Option<String> {
    pkg.targets.iter()
               .find(|target| target.kind.iter().any(|k| k == "lib"))
               .map(|target| format!("extern crate {};\n", target.name.replace("-", "_")))
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File, pkg: &Package) -> Result<(), Box<error::Error>> {
write!(script, r#"#![no_main]
extern crate libfuzzer_sys;
{}
#[export_name="rust_fuzzer_test_input"]
pub extern fn go(data: &[u8]) {{
    // fuzzed code goes here
}}"#, link_to_lib(pkg).unwrap_or(String::new())).map_err(|e| e.into())
}

/// Add a new fuzz target script with a given name
fn add_target(target: String) -> Result<(), Box<error::Error>> {
    let target_file = format!("fuzz/fuzzers/{}.rs", target);
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    let me = get_package();
    dummy_target(&mut script, &me)?;

    let mut cargo = fs::OpenOptions::new().append(true).open(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[[bin]]
name = "{0}"
path = "fuzzers/{0}.rs"
"#, target).map_err(|e| e.into())

}

/// Build or rebuild libFuzzer (rebuilds only if the compiler version changed)
///
/// We can't just use libFuzzer as a dependency since libgcc will
/// get compiled with sanitizer support. RUSTFLAGS does not discriminate
/// between build dependencies and regular ones.
///
/// https://github.com/rust-lang/cargo/issues/3739
fn rebuild_libfuzzer() -> Result<(), Box<error::Error>> {
    if let Err(_) = env::set_current_dir("./libfuzzer") {
        let mut git = process::Command::new("git");
        let mut cmd = git.arg("clone")
                         .arg("https://github.com/rust-fuzz/libfuzzer-sys.git")
                         .arg("libfuzzer");
        let result = cmd.spawn()?.wait()?;
        if !result.success() {
            return Err("Failed to clone libfuzzer-sys".into())
        }
        env::set_current_dir("./libfuzzer")?;
    }
    let mut cmd = process::Command::new("cargo");
    cmd.arg("build")
       .arg("--release")
       .spawn()?
       .wait()?;

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build libfuzzer-sys".into())
    }
    env::set_current_dir("..")
        .map_err(|e| e.into())
}

/// Fuzz a given fuzz target
fn run_target(target: String) -> Result<bool, Box<error::Error>> {
    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
    if !flags.is_empty() {
        flags.push(' ');
    }
    flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS", &flags);

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build fuzz target".into())
    }

    if let Err(k) = fs::create_dir("corpus") {
        if k.kind() == io::ErrorKind::AlreadyExists {
            // do nothing
        } else {
            return Err(k.into());
        }
    }

    // can't use cargo run since we can't pass -L args to it
    let path = format!("target/debug/{}", target);
    let mut run_cmd = process::Command::new(path);
    run_cmd.arg("corpus")
           .env("ASAN_OPTIONS", "detect_odr_violation=0");
    let result = run_cmd.spawn()?.wait()?;
    Ok(result.success())
}
<commit_msg>Support artifact directory (fix #6)<commit_after>// Copyright 2016 rust-fuzz developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

extern crate cargo_metadata;
extern crate docopt;
extern crate rustc_serialize;

use cargo_metadata::{metadata, Package};
use docopt::Docopt;
use std::{env, error, fs, io, path, process};
use std::io::Write;

const USAGE: &'static str = "
Cargo Fuzz

Usage:
  cargo fuzz --init
  cargo fuzz --fuzz-target TARGET
  cargo fuzz --add TARGET
  cargo fuzz --list
  cargo fuzz (-h | --help)

Options:
  -h --help              Show this screen.
  --init                 Initialize fuzz folder
  --fuzz-target TARGET   Run with given fuzz target in fuzz/fuzzers
  --add TARGET           Add a new fuzz target
  --list                 List the available fuzz targets
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_init: bool,
    flag_add: Option<String>,
    flag_fuzz_target: Option<String>,
    flag_list: bool,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let result = if args.flag_init {
        init_fuzz()
    } else if let Some(target) = args.flag_add {
        add_target(target)
    } else if let Some(target) = args.flag_fuzz_target {
        let result = run_target(target);
        if let Ok(success) = result {
            if success {
                // Can this ever happen?
                Ok(())
            } else {
                println!("Fuzzing found errors!");
                process::exit(-1)
            }
        } else {
            result.map(|_| ())
        }
    } else if args.flag_list {
        list_fuzz_targets()
            .map(|_| ())
    } else {
        println!("Invalid arguments. Usage:\n{}", USAGE);
        return;
    };
    if let Err(error) = result {
        writeln!(io::stderr(), "Error: {}", error)
            .expect("failed writing to stderr");
    }
}

fn list_fuzz_targets() -> Result<(), Box<error::Error>> {
    if !path::Path::new("./fuzz").is_dir() {
        return Err("Fuzzing crate has not been initialized. Run `cargo fuzz --init` to initialize it.".into());
    }
    env::set_current_dir("./fuzz")?;
    let package = get_package();
    for target in &package.targets {
        println!("{}", target.name);
    }
    Ok(())
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> Result<(), Box<error::Error>> {
    let me = get_package();

    fs::create_dir("./fuzz")?;
    fs::create_dir("./fuzz/fuzzers")?;

    let mut cargo = fs::File::create(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[package]
name = "{0}-fuzz"
version = "0.0.1"
authors = ["Automatically generated"]
publish = false

[dependencies.{0}]
path = ".."

# Prevent this from interfering with workspaces
[workspace]
members = ["."]

[[bin]]
name = "fuzzer_script_1"
path = "fuzzers/fuzzer_script_1.rs"
"#, me.name)?;

    let mut ignore = fs::File::create(path::Path::new("./fuzz/.gitignore"))?;

write!(ignore, r#"
target
libfuzzer
corpus
artifacts
"#)?;

    let mut script = fs::File::create(path::Path::new("./fuzz/fuzzers/fuzzer_script_1.rs"))?;
    dummy_target(&mut script, &me)
}

/// Returns metadata for the Cargo package in the current directory
fn get_package() -> Package {
    // todo error handling
    let meta = metadata(None).unwrap();
    let mut p = env::current_dir().unwrap();
    p.push("Cargo.toml");
    let p = p.to_str().unwrap();
    meta.packages.into_iter().find(|package| package.manifest_path == p).unwrap()
}

/// If the package contains a library target, generate an `extern crate` line to link to it.
fn link_to_lib(pkg: &Package) -> Option<String> {
    pkg.targets.iter()
               .find(|target| target.kind.iter().any(|k| k == "lib"))
               .map(|target| format!("extern crate {};\n", target.name.replace("-", "_")))
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File, pkg: &Package) -> Result<(), Box<error::Error>> {
write!(script, r#"#![no_main]
extern crate libfuzzer_sys;
{}
#[export_name="rust_fuzzer_test_input"]
pub extern fn go(data: &[u8]) {{
    // fuzzed code goes here
}}"#, link_to_lib(pkg).unwrap_or(String::new())).map_err(|e| e.into())
}

/// Add a new fuzz target script with a given name
fn add_target(target: String) -> Result<(), Box<error::Error>> {
    let target_file = format!("fuzz/fuzzers/{}.rs", target);
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    let me = get_package();
    dummy_target(&mut script, &me)?;

    let mut cargo = fs::OpenOptions::new().append(true).open(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[[bin]]
name = "{0}"
path = "fuzzers/{0}.rs"
"#, target).map_err(|e| e.into())

}

/// Build or rebuild libFuzzer (rebuilds only if the compiler version changed)
///
/// We can't just use libFuzzer as a dependency since libgcc will
/// get compiled with sanitizer support. RUSTFLAGS does not discriminate
/// between build dependencies and regular ones.
///
/// https://github.com/rust-lang/cargo/issues/3739
fn rebuild_libfuzzer() -> Result<(), Box<error::Error>> {
    if let Err(_) = env::set_current_dir("./libfuzzer") {
        let mut git = process::Command::new("git");
        let mut cmd = git.arg("clone")
                         .arg("https://github.com/rust-fuzz/libfuzzer-sys.git")
                         .arg("libfuzzer");
        let result = cmd.spawn()?.wait()?;
        if !result.success() {
            return Err("Failed to clone libfuzzer-sys".into())
        }
        env::set_current_dir("./libfuzzer")?;
    }
    let mut cmd = process::Command::new("cargo");
    cmd.arg("build")
       .arg("--release")
       .spawn()?
       .wait()?;

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build libfuzzer-sys".into())
    }
    env::set_current_dir("..")
        .map_err(|e| e.into())
}

fn make_dir_if_not_exist(dir: &str) -> Result<(), io::Error> {
    if let Err(k) = fs::create_dir(dir) {
        if k.kind() == io::ErrorKind::AlreadyExists {
            // do nothing
        } else {
            return Err(k);
        }
    }
    Ok(())
}
/// Fuzz a given fuzz target
fn run_target(target: String) -> Result<bool, Box<error::Error>> {
    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
    if !flags.is_empty() {
        flags.push(' ');
    }
    flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS", &flags);

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build fuzz target".into())
    }

    make_dir_if_not_exist("corpus")?;
    make_dir_if_not_exist("artifacts")?;

    // can't use cargo run since we can't pass -L args to it
    let path = format!("target/debug/{}", target);
    let mut run_cmd = process::Command::new(path);
    run_cmd.arg("-artifact_prefix=artifacts/")
           .arg("corpus") // must be last arg
           .env("ASAN_OPTIONS", "detect_odr_violation=0");
    let result = run_cmd.spawn()?.wait()?;
    Ok(result.success())
}
<|endoftext|>


<commit_before>use std::{fmt, ops::Bound};

use crate::Lazy;

use super::*;

#[derive(Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct Node {
    pub(crate) data: Data,
    pub(crate) next: Option<PageId>,
    pub(crate) lo: IVec,
    pub(crate) hi: IVec,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
}

impl Debug for Node {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let data = self.data.fmt_keys(&self.lo);

        write!(
            f,
            "Node {{ \
             lo: {:?} \
             hi: {:?} \
             next: {:?} \
             merging_child: {:?} \
             merging: {} \
             data: {:?} }}",
            self.lo, self.hi, self.next, self.merging_child, self.merging, data
        )
    }
}

impl Node {
    pub(crate) fn apply(&mut self, frag: &Frag) {
        use self::Frag::*;

        assert!(
            !self.merging,
            "somehow a frag was applied to a node after it was merged"
        );

        match *frag {
            Set(ref k, ref v) => {
                // (when hi is empty, it means it's unbounded)
                if self.hi.is_empty()
                    || prefix_cmp_encoded(k, &self.hi, &self.lo)
                        == std::cmp::Ordering::Less
                {
                    self.set_leaf(k.clone(), v.clone());
                } else {
                    panic!(
                        "tried to consolidate set at key <= hi.\
                         Set({:?}, {:?}) to node {:?}",
                        k, v, self
                    )
                }
            }
            Del(ref k) => {
                // (when hi is empty, it means it's unbounded)
                if self.hi.is_empty()
                    || prefix_cmp_encoded(k, &self.hi, &self.lo)
                        == std::cmp::Ordering::Less
                {
                    self.del_leaf(k);
                } else {
                    panic!("tried to consolidate del at key <= hi")
                }
            }
            Base(_) => panic!("trying to apply a Base to frag {:?}", self),
            ParentMergeIntention(pid) => {
                assert!(
                    self.merging_child.is_none(),
                    "trying to merge {:?} into node {:?} which \
                     is already merging another child",
                    frag,
                    self
                );
                self.merging_child = Some(pid);
            }
            ParentMergeConfirm => {
                assert!(self.merging_child.is_some());
                let merged_child = self.merging_child.take().expect(
                    "we should have a specific \
                     child that was merged if this \
                     frag appears here",
                );
                self.data.parent_merge_confirm(merged_child);
            }
            ChildMergeCap => {
                self.merging = true;
            }
        }
    }

    pub(crate) fn set_leaf(&mut self, key: IVec, val: IVec) {
        if let Data::Leaf(ref mut records) = self.data {
            let search = records.binary_search_by(|(k, _)| prefix_cmp(k, &key));
            match search {
                Ok(idx) => records[idx] = (key, val),
                Err(idx) => records.insert(idx, (key, val)),
            }
        } else {
            panic!("tried to Set a value to an index");
        }
    }

    pub(crate) fn parent_split(&mut self, at: &[u8], to: PageId) -> bool {
        if let Data::Index(ref mut ptrs) = self.data {
            let encoded_sep = prefix_encode(&self.lo, at);
            match ptrs.binary_search_by(|a| prefix_cmp(&a.0, &encoded_sep)) {
                Ok(_) => {
                    debug!(
                        "parent_split skipped because \
                         parent already contains child at split point \
                         due to deep race"
                    );
                    return false;
                }
                Err(idx) => ptrs.insert(idx, (encoded_sep, to)),
            }
        } else {
            panic!("tried to attach a ParentSplit to a Leaf chain");
        }

        true
    }

    pub(crate) fn del_leaf(&mut self, key: &IVec) {
        if let Data::Leaf(ref mut records) = self.data {
            let search = records
                .binary_search_by(|&(ref k, ref _v)| prefix_cmp(k, &*key));
            if let Ok(idx) = search {
                let _removed = records.remove(idx);
            }
        } else {
            panic!("tried to attach a Del to an Index chain");
        }
    }

    pub(crate) fn split(mut self) -> (Self, Self) {
        let (split, right_data) = self.data.split(&self.lo);
        let rhs = Self {
            data: right_data,
            next: self.next,
            lo: split,
            hi: self.hi.clone(),
            merging_child: None,
            merging: false,
        };

        self.data.drop_gte(&rhs.lo, &self.lo);
        self.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        self.next = None;

        (self, rhs)
    }

    pub(crate) fn receive_merge(&self, rhs: &Self) -> Self {
        let mut merged = self.clone();
        merged.hi = rhs.hi.clone();
        merged.data.receive_merge(
            rhs.lo.as_ref(),
            merged.lo.as_ref(),
            &rhs.data,
        );
        merged.next = rhs.next;
        merged
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Excluded(bound) if self.hi >= *bound => true,
            Bound::Included(bound) if self.hi > *bound => true,
            _ => self.hi.is_empty(),
        }
    }

    pub(crate) fn contains_lower_bound(
        &self,
        bound: &Bound<IVec>,
        is_forward: bool,
    ) -> bool {
        match bound {
            Bound::Excluded(bound)
                if self.lo < *bound || (is_forward && *bound == self.lo) =>
            {
                true
            }
            Bound::Included(bound) if self.lo <= *bound => true,
            Bound::Unbounded if !is_forward => self.hi.is_empty(),
            _ => self.lo.is_empty(),
        }
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.data.is_index());

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(&self.lo, &self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, &self.lo);
                prefix_encode(&self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, &self.lo);
                prefix_encode(&self.lo, max)
            }
        };

        let records = self.data.leaf_ref().unwrap();
        let search =
            records.binary_search_by(|(k, _)| prefix_cmp(k, &predecessor_key));

        let idx = match search {
            Ok(idx) => idx,
            Err(idx) if idx < records.len() => idx,
            _ => return None,
        };

        for (k, v) in &records[idx..] {
            match bound {
                Bound::Excluded(b)
                    if prefix_cmp_encoded(k, b, &self.lo)
                        == std::cmp::Ordering::Equal =>
                {
                    // keep going because we wanted to exclude
                    // this key.
                    continue;
                }
                _ => {}
            }
            let decoded_key = prefix_decode(&self.lo, &k);
            return Some((IVec::from(decoded_key), v.clone()));
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
    ) -> Option<(IVec, IVec)> {
        static MAX_IVEC: Lazy<IVec, fn() -> IVec> = Lazy::new(init_max_ivec);

        fn init_max_ivec() -> IVec {
            let mut base = vec![255; 1024 * 1024];
            base[0] = 0;
            IVec::from(base)
        }

        assert!(!self.data.is_index());

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    MAX_IVEC.clone()
                } else {
                    prefix_encode(&self.lo, &self.hi)
                }
            }
            Bound::Included(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, &self.hi)
                };
                prefix_encode(&self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, &self.hi)
                };
                prefix_encode(&self.lo, min)
            }
        };

        let records = self.data.leaf_ref().unwrap();
        let search =
            records.binary_search_by(|(k, _)| prefix_cmp(k, &successor_key));

        let idx = match search {
            Ok(idx) => idx,
            Err(idx) if idx > 0 => idx - 1,
            _ => return None,
        };

        for (k, v) in records[0..=idx].iter().rev() {
            match bound {
                Bound::Excluded(b)
                    if prefix_cmp_encoded(k, b, &self.lo)
                        == std::cmp::Ordering::Equal =>
                {
                    // keep going because we wanted to exclude
                    // this key.
                    continue;
                }
                _ => {}
            }
            let decoded_key = prefix_decode(&self.lo, &k);
            return Some((IVec::from(decoded_key), v.clone()));
        }
        None
    }

    /// leaf_pair_for_key finds an existing value pair for a given key.
    pub(crate) fn leaf_pair_for_key(
        &self,
        key: &[u8],
    ) -> Option<(&IVec, &IVec)> {
        assert!(!self.data.is_index());

        let records = self.data.leaf_ref().unwrap();

        let common_prefix = key
            .iter()
            .zip(self.lo.iter())
            .take(u8::max_value() as usize)
            .take_while(|(a, b)| a == b)
            .count() as u8;

        let search = records
            .binary_search_by(|&(ref k, ref _v)| {
                if k[0] > common_prefix {
                    Ordering::Less
                } else if k[0] < common_prefix {
                    Ordering::Greater
                } else {
                    k[1..].cmp(&key[common_prefix as usize..])
                }
            })
            .ok();

        search.map(|idx| (&records[idx].0, &records[idx].1))
    }

    /// node_kv_pair returns either existing (node/key, value) pair or
    /// (node/key, none) where a node/key is node level encoded key.
    pub fn node_kv_pair(&self, key: &[u8]) -> (IVec, Option<IVec>) {
        if let Some((k, v)) = self.leaf_pair_for_key(key.as_ref()) {
            (k.clone(), Some(v.clone()))
        } else {
            let encoded_key = prefix_encode(&self.lo, key.as_ref());
            let encoded_val = None;
            (encoded_key, encoded_val)
        }
    }

    pub(crate) fn should_split(&self) -> bool {
        let threshold = if cfg!(any(test, feature = "lock_free_delays")) {
            2
        } else if self.data.is_index() {
            256
        } else {
            16
        };

        let size_checks = self.data.len() > threshold;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn should_merge(&self) -> bool {
        let threshold = if cfg!(any(test, feature = "lock_free_delays")) {
            1
        } else if self.data.is_index() {
            64
        } else {
            4
        };

        let size_checks = self.data.len() < threshold;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> (usize, PageId) {
        let records = self
            .data
            .index_ref()
            .expect("index_next_node called on leaf");

        let common_prefix = key
            .iter()
            .zip(self.lo.iter())
            .take(u8::max_value() as usize)
            .take_while(|(a, b)| a == b)
            .count() as u8;

        let search = binary_search_lub(records, |&(ref k, ref _v)| {
            if k[0] > common_prefix {
                Ordering::Less
            } else if k[0] < common_prefix {
                Ordering::Greater
            } else {
                k[1..].cmp(&key[common_prefix as usize..])
            }
        });

        let index = search.expect("failed to traverse index");

        (index, records[index].1)
    }
}
<commit_msg>Remove redundant assertion<commit_after>use std::{fmt, ops::Bound};

use crate::Lazy;

use super::*;

#[derive(Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct Node {
    pub(crate) data: Data,
    pub(crate) next: Option<PageId>,
    pub(crate) lo: IVec,
    pub(crate) hi: IVec,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
}

impl Debug for Node {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let data = self.data.fmt_keys(&self.lo);

        write!(
            f,
            "Node {{ \
             lo: {:?} \
             hi: {:?} \
             next: {:?} \
             merging_child: {:?} \
             merging: {} \
             data: {:?} }}",
            self.lo, self.hi, self.next, self.merging_child, self.merging, data
        )
    }
}

impl Node {
    pub(crate) fn apply(&mut self, frag: &Frag) {
        use self::Frag::*;

        assert!(
            !self.merging,
            "somehow a frag was applied to a node after it was merged"
        );

        match *frag {
            Set(ref k, ref v) => {
                // (when hi is empty, it means it's unbounded)
                if self.hi.is_empty()
                    || prefix_cmp_encoded(k, &self.hi, &self.lo)
                        == std::cmp::Ordering::Less
                {
                    self.set_leaf(k.clone(), v.clone());
                } else {
                    panic!(
                        "tried to consolidate set at key <= hi.\
                         Set({:?}, {:?}) to node {:?}",
                        k, v, self
                    )
                }
            }
            Del(ref k) => {
                // (when hi is empty, it means it's unbounded)
                if self.hi.is_empty()
                    || prefix_cmp_encoded(k, &self.hi, &self.lo)
                        == std::cmp::Ordering::Less
                {
                    self.del_leaf(k);
                } else {
                    panic!("tried to consolidate del at key <= hi")
                }
            }
            Base(_) => panic!("trying to apply a Base to frag {:?}", self),
            ParentMergeIntention(pid) => {
                assert!(
                    self.merging_child.is_none(),
                    "trying to merge {:?} into node {:?} which \
                     is already merging another child",
                    frag,
                    self
                );
                self.merging_child = Some(pid);
            }
            ParentMergeConfirm => {
                assert!(self.merging_child.is_some());
                let merged_child = self.merging_child.take().expect(
                    "we should have a specific \
                     child that was merged if this \
                     frag appears here",
                );
                self.data.parent_merge_confirm(merged_child);
            }
            ChildMergeCap => {
                self.merging = true;
            }
        }
    }

    pub(crate) fn set_leaf(&mut self, key: IVec, val: IVec) {
        if let Data::Leaf(ref mut records) = self.data {
            let search = records.binary_search_by(|(k, _)| prefix_cmp(k, &key));
            match search {
                Ok(idx) => records[idx] = (key, val),
                Err(idx) => records.insert(idx, (key, val)),
            }
        } else {
            panic!("tried to Set a value to an index");
        }
    }

    pub(crate) fn parent_split(&mut self, at: &[u8], to: PageId) -> bool {
        if let Data::Index(ref mut ptrs) = self.data {
            let encoded_sep = prefix_encode(&self.lo, at);
            match ptrs.binary_search_by(|a| prefix_cmp(&a.0, &encoded_sep)) {
                Ok(_) => {
                    debug!(
                        "parent_split skipped because \
                         parent already contains child at split point \
                         due to deep race"
                    );
                    return false;
                }
                Err(idx) => ptrs.insert(idx, (encoded_sep, to)),
            }
        } else {
            panic!("tried to attach a ParentSplit to a Leaf chain");
        }

        true
    }

    pub(crate) fn del_leaf(&mut self, key: &IVec) {
        if let Data::Leaf(ref mut records) = self.data {
            let search = records
                .binary_search_by(|&(ref k, ref _v)| prefix_cmp(k, &*key));
            if let Ok(idx) = search {
                let _removed = records.remove(idx);
            }
        } else {
            panic!("tried to attach a Del to an Index chain");
        }
    }

    pub(crate) fn split(mut self) -> (Self, Self) {
        let (split, right_data) = self.data.split(&self.lo);
        let rhs = Self {
            data: right_data,
            next: self.next,
            lo: split,
            hi: self.hi.clone(),
            merging_child: None,
            merging: false,
        };

        self.data.drop_gte(&rhs.lo, &self.lo);
        self.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        self.next = None;

        (self, rhs)
    }

    pub(crate) fn receive_merge(&self, rhs: &Self) -> Self {
        let mut merged = self.clone();
        merged.hi = rhs.hi.clone();
        merged.data.receive_merge(
            rhs.lo.as_ref(),
            merged.lo.as_ref(),
            &rhs.data,
        );
        merged.next = rhs.next;
        merged
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Excluded(bound) if self.hi >= *bound => true,
            Bound::Included(bound) if self.hi > *bound => true,
            _ => self.hi.is_empty(),
        }
    }

    pub(crate) fn contains_lower_bound(
        &self,
        bound: &Bound<IVec>,
        is_forward: bool,
    ) -> bool {
        match bound {
            Bound::Excluded(bound)
                if self.lo < *bound || (is_forward && *bound == self.lo) =>
            {
                true
            }
            Bound::Included(bound) if self.lo <= *bound => true,
            Bound::Unbounded if !is_forward => self.hi.is_empty(),
            _ => self.lo.is_empty(),
        }
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.data.is_index());

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(&self.lo, &self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, &self.lo);
                prefix_encode(&self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, &self.lo);
                prefix_encode(&self.lo, max)
            }
        };

        let records = self.data.leaf_ref().unwrap();
        let search =
            records.binary_search_by(|(k, _)| prefix_cmp(k, &predecessor_key));

        let idx = match search {
            Ok(idx) => idx,
            Err(idx) if idx < records.len() => idx,
            _ => return None,
        };

        for (k, v) in &records[idx..] {
            match bound {
                Bound::Excluded(b)
                    if prefix_cmp_encoded(k, b, &self.lo)
                        == std::cmp::Ordering::Equal =>
                {
                    // keep going because we wanted to exclude
                    // this key.
                    continue;
                }
                _ => {}
            }
            let decoded_key = prefix_decode(&self.lo, &k);
            return Some((IVec::from(decoded_key), v.clone()));
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
    ) -> Option<(IVec, IVec)> {
        static MAX_IVEC: Lazy<IVec, fn() -> IVec> = Lazy::new(init_max_ivec);

        fn init_max_ivec() -> IVec {
            let mut base = vec![255; 1024 * 1024];
            base[0] = 0;
            IVec::from(base)
        }

        assert!(!self.data.is_index());

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    MAX_IVEC.clone()
                } else {
                    prefix_encode(&self.lo, &self.hi)
                }
            }
            Bound::Included(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, &self.hi)
                };
                prefix_encode(&self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, &self.hi)
                };
                prefix_encode(&self.lo, min)
            }
        };

        let records = self.data.leaf_ref().unwrap();
        let search =
            records.binary_search_by(|(k, _)| prefix_cmp(k, &successor_key));

        let idx = match search {
            Ok(idx) => idx,
            Err(idx) if idx > 0 => idx - 1,
            _ => return None,
        };

        for (k, v) in records[0..=idx].iter().rev() {
            match bound {
                Bound::Excluded(b)
                    if prefix_cmp_encoded(k, b, &self.lo)
                        == std::cmp::Ordering::Equal =>
                {
                    // keep going because we wanted to exclude
                    // this key.
                    continue;
                }
                _ => {}
            }
            let decoded_key = prefix_decode(&self.lo, &k);
            return Some((IVec::from(decoded_key), v.clone()));
        }
        None
    }

    /// leaf_pair_for_key finds an existing value pair for a given key.
    pub(crate) fn leaf_pair_for_key(
        &self,
        key: &[u8],
    ) -> Option<(&IVec, &IVec)> {
        let records = self
            .data
            .leaf_ref()
            .expect("leaf_pair_for_key called on index node");

        let common_prefix = key
            .iter()
            .zip(self.lo.iter())
            .take(u8::max_value() as usize)
            .take_while(|(a, b)| a == b)
            .count() as u8;

        let search = records
            .binary_search_by(|&(ref k, ref _v)| {
                if k[0] > common_prefix {
                    Ordering::Less
                } else if k[0] < common_prefix {
                    Ordering::Greater
                } else {
                    k[1..].cmp(&key[common_prefix as usize..])
                }
            })
            .ok();

        search.map(|idx| (&records[idx].0, &records[idx].1))
    }

    /// node_kv_pair returns either existing (node/key, value) pair or
    /// (node/key, none) where a node/key is node level encoded key.
    pub fn node_kv_pair(&self, key: &[u8]) -> (IVec, Option<IVec>) {
        if let Some((k, v)) = self.leaf_pair_for_key(key.as_ref()) {
            (k.clone(), Some(v.clone()))
        } else {
            let encoded_key = prefix_encode(&self.lo, key.as_ref());
            let encoded_val = None;
            (encoded_key, encoded_val)
        }
    }

    pub(crate) fn should_split(&self) -> bool {
        let threshold = if cfg!(any(test, feature = "lock_free_delays")) {
            2
        } else if self.data.is_index() {
            256
        } else {
            16
        };

        let size_checks = self.data.len() > threshold;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn should_merge(&self) -> bool {
        let threshold = if cfg!(any(test, feature = "lock_free_delays")) {
            1
        } else if self.data.is_index() {
            64
        } else {
            4
        };

        let size_checks = self.data.len() < threshold;
        let safety_checks = self.merging_child.is_none() && !self.merging;

        size_checks && safety_checks
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> (usize, PageId) {
        let records = self
            .data
            .index_ref()
            .expect("index_next_node called on leaf");

        let common_prefix = key
            .iter()
            .zip(self.lo.iter())
            .take(u8::max_value() as usize)
            .take_while(|(a, b)| a == b)
            .count() as u8;

        let search = binary_search_lub(records, |&(ref k, ref _v)| {
            if k[0] > common_prefix {
                Ordering::Less
            } else if k[0] < common_prefix {
                Ordering::Greater
            } else {
                k[1..].cmp(&key[common_prefix as usize..])
            }
        });

        let index = search.expect("failed to traverse index");

        (index, records[index].1)
    }
}
<|endoftext|>


<commit_before>//! Contains all HTML-related code. See the `Page` documentation for details.

// TODO: Why do I need these two lines "extern crate html5ever" and "extern crate
// html5ever_dom_sink" both here and in main.rs?
extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate rand;
extern crate tendril;

use std::str::FromStr;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use regex::Captures;
use regex::Regex;

use wiki::Wiki;

use ::START_MARKER;
use ::END_MARKER;

// TODO: massive cleanup, all over this file.

/// Represents, and owns all behavior related to, the contents of the HTML page shown to the
/// user. This includes fetching the rendered article from Wikipedia, replacing its contents with
/// the rendered wikitext, and processing/removing merge markers.
///
/// The API to this struct comprises two methods: Page::new() constructs a new Page. It should be
/// called early, so the Page can start fetching article HTML from Wikipedia.
/// Page:replace_body_and_remove_merge_markers() processes the merge markers in the rendered
/// wikitext and puts the header and footer around it.
pub struct Page {
    /// The string used as a placeholder for the article body in the page skeleton.
    placeholder: String,
    /// The Receiver that will receive the page skeleton when it's been fetched and processed.
    page_skeleton_receiver: Receiver<Result<String, String>>,
}

impl Page {
    /// Creates a new Page representing the article at `title`. This kicks off a background thread
    /// that fetches the current article HTML from Wikipedia. Because of that, it should be called
    /// as early as possible (as soon as the title being served is known), so that the page fetch
    /// stays off the critical path for page load.
    pub fn new(title: &str, wiki: Wiki) -> Page {
        let placeholder = format!("WMW_PLACEHOLDER_{}", rand::random::<u64>());
        let page_skeleton_receiver =
            Page::spawn_page_skeleton_fetch_thread(title, placeholder.clone(), wiki);
        Page {
            placeholder: placeholder,
            page_skeleton_receiver: page_skeleton_receiver,
        }
    }

    /// This finishes the HTML processing - it replaces the merge markers in `article_body` with
    /// HTML tags, and inserts the resulting HTML into the page skeleton.
    pub fn replace_body_and_remove_merge_markers(&self, article_body: String)
                                                 -> Result<String, String> {
        match self.page_skeleton_receiver.recv() {
            Ok(Ok(page_skeleton)) => {
                let finished_article_body = process_merge_markers(article_body);
                Ok(page_skeleton.replace(&self.placeholder, &finished_article_body))
            },
            Ok(Err(msg))=> Err(msg),
            Err(err) => Err(format!("error: {}", err)),
        }
    }

    fn spawn_page_skeleton_fetch_thread(title: &str, placeholder: String, wiki: Wiki)
                                  -> Receiver<Result<String, String>> {
        let (page_skeleton_sender, page_skeleton_receiver) = channel::<Result<String, String>>();
        let title = title.to_owned().clone();
        thread::spawn(move|| {
            page_skeleton_sender.send(
                match wiki.get_current_page_content(&title) {
                    Ok(content) =>
                        replace_node_with_placeholder(&content, "mw-content-text", &placeholder),
                    Err(msg) => Err(msg),
                }).unwrap();
        });
        page_skeleton_receiver
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str)
    -> Result<String, String> {
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    try_display!(
        html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default()),
        "Failed to serialize modified HTML");
    Ok(try_display!(String::from_utf8(serialized),
                    "Error converting serialized HTML to UTF-8 string"))
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
        return attributes.into_iter().any(
            |attribute| attribute.name.local.as_slice() == "id" &&
                format!("{}", attribute.value) == id);
    }

    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => (&node.children).into_iter()
            .map(|child| find_node_by_id(child, id))
            .filter(|result| result.is_ok())
            .map(|result| result.unwrap())
            .next().ok_or(format!("No node with ID {} found", id)),
    }
}

/// Removes merge markers that are inside HTML tags, and replaces the others with  tags.
fn process_merge_markers(html: String) -> String {
    let start_regex = Regex::new(&format!("{}([0-9]+){}", START_MARKER, START_MARKER)).unwrap();
    let end_regex = Regex::new(&format!("{}[0-9]+{}", END_MARKER, END_MARKER)).unwrap();

    let html = remove_merge_markers(html);
    let html = start_regex.replace_all(
        &html, |captures: &Captures| format!("",
                                             captures.at(1).unwrap()));
    end_regex.replace_all(&html, "")
}

fn remove_merge_markers(html: String) -> String {
    // Finds markers where the end, but not the start, is inside a tag.
    let regex1 = Regex::new(&format!(
        r"{}[0-9]+{}([^{}]*?)<([^>]*?){}[0-9]+{}([^>]*?)>",
        START_MARKER, START_MARKER, END_MARKER, END_MARKER, END_MARKER)).unwrap();
    // Finds markers where the start, but not the end, is inside a tag.
    let regex2 = Regex::new(&format!(
        r"<([^>]*?){}[0-9]+{}([^>]*?)>([^{}]*?){}[0-9]+{}",
        START_MARKER, START_MARKER, END_MARKER, END_MARKER, END_MARKER)).unwrap();
    // Finds markers where both the start and end are inside tags.
    let regex3 = Regex::new(&format!(
        r"<([^>]*?){}[0-9]+{}([^>]*?)>([^{}{}]*?)<([^>]*?){}[0-9]+{}([^>]*?)>",
        START_MARKER, START_MARKER, START_MARKER, END_MARKER, END_MARKER, END_MARKER)).unwrap();
    let html = regex1.replace_all(
        &html, |captures: &Captures|
        format!("{}<{}{}>", captures.at(1).unwrap(), captures.at(2).unwrap(),
                captures.at(3).unwrap()));
    let html = regex2.replace_all(
        &html, |captures: &Captures|
        format!("<{}{}>{}", captures.at(1).unwrap(), captures.at(2).unwrap(),
                captures.at(3).unwrap()));
    regex3.replace_all(
        &html, |captures: &Captures|
        format!("<{}{}>{}<{}{}>", captures.at(1).unwrap(), captures.at(2).unwrap(),
                captures.at(3).unwrap(), captures.at(4).unwrap(), captures.at(5).unwrap()))
}

#[cfg(test)]
mod tests {
    use super::{remove_merge_markers, replace_node_with_placeholder};
    use ::START_MARKER;
    use ::END_MARKER;

    fn test_process_merge_markers() {
        let html = format!(
            "{}456{}{}456{}{}123{}t",
            START_MARKER, START_MARKER, END_MARKER, END_MARKER,
            START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected_regex =
            regex!("]*>");
        assert!(expected_regex.is_match(&html));
    }

    #[test]
    fn test_remove_merge_markers_keep() {
        let html = format!("{}456{}{}456{}",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        assert_eq!(html.clone(), remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_keep_one_remove_one() {
        let html = format!(
            "{}234{}text{}234{}{}567{}",
            START_MARKER, START_MARKER, END_MARKER, END_MARKER, START_MARKER, START_MARKER,
            END_MARKER, END_MARKER);
        let expected = format!(
            "{}234{}text{}234{}",
            START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_end_inside_tag() {
        let html = format!("{}123{}",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected = "";
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_start_inside_tag() {
        let html = format!("{}123{}",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected = "";
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_both_inside_tag() {
        let html = format!("text",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected = "text";
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<commit_msg>Remove a "general cleanup" TODO from page.rs, which has now been cleaned up.<commit_after>//! Contains all HTML-related code. See the `Page` documentation for details.

// TODO: Why do I need these two lines "extern crate html5ever" and "extern crate
// html5ever_dom_sink" both here and in main.rs?
extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate rand;
extern crate tendril;

use std::str::FromStr;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use regex::Captures;
use regex::Regex;

use wiki::Wiki;

use ::START_MARKER;
use ::END_MARKER;

/// Represents, and owns all behavior related to, the contents of the HTML page shown to the
/// user. This includes fetching the rendered article from Wikipedia, replacing its contents with
/// the rendered wikitext, and processing/removing merge markers.
///
/// The API to this struct comprises two methods: Page::new() constructs a new Page. It should be
/// called early, so the Page can start fetching article HTML from Wikipedia.
/// Page:replace_body_and_remove_merge_markers() processes the merge markers in the rendered
/// wikitext and puts the header and footer around it.
pub struct Page {
    /// The string used as a placeholder for the article body in the page skeleton.
    placeholder: String,
    /// The Receiver that will receive the page skeleton when it's been fetched and processed.
    page_skeleton_receiver: Receiver<Result<String, String>>,
}

impl Page {
    /// Creates a new Page representing the article at `title`. This kicks off a background thread
    /// that fetches the current article HTML from Wikipedia. Because of that, it should be called
    /// as early as possible (as soon as the title being served is known), so that the page fetch
    /// stays off the critical path for page load.
    pub fn new(title: &str, wiki: Wiki) -> Page {
        let placeholder = format!("WMW_PLACEHOLDER_{}", rand::random::<u64>());
        let page_skeleton_receiver =
            Page::spawn_page_skeleton_fetch_thread(title, placeholder.clone(), wiki);
        Page {
            placeholder: placeholder,
            page_skeleton_receiver: page_skeleton_receiver,
        }
    }

    /// This finishes the HTML processing - it replaces the merge markers in `article_body` with
    /// HTML tags, and inserts the resulting HTML into the page skeleton.
    pub fn replace_body_and_remove_merge_markers(&self, article_body: String)
                                                 -> Result<String, String> {
        match self.page_skeleton_receiver.recv() {
            Ok(Ok(page_skeleton)) => {
                let finished_article_body = process_merge_markers(article_body);
                Ok(page_skeleton.replace(&self.placeholder, &finished_article_body))
            },
            Ok(Err(msg))=> Err(msg),
            Err(err) => Err(format!("error: {}", err)),
        }
    }

    fn spawn_page_skeleton_fetch_thread(title: &str, placeholder: String, wiki: Wiki)
                                  -> Receiver<Result<String, String>> {
        let (page_skeleton_sender, page_skeleton_receiver) = channel::<Result<String, String>>();
        let title = title.to_owned().clone();
        thread::spawn(move|| {
            page_skeleton_sender.send(
                match wiki.get_current_page_content(&title) {
                    Ok(content) =>
                        replace_node_with_placeholder(&content, "mw-content-text", &placeholder),
                    Err(msg) => Err(msg),
                }).unwrap();
        });
        page_skeleton_receiver
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str)
    -> Result<String, String> {
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    try_display!(
        html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default()),
        "Failed to serialize modified HTML");
    Ok(try_display!(String::from_utf8(serialized),
                    "Error converting serialized HTML to UTF-8 string"))
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
        return attributes.into_iter().any(
            |attribute| attribute.name.local.as_slice() == "id" &&
                format!("{}", attribute.value) == id);
    }

    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => (&node.children).into_iter()
            .map(|child| find_node_by_id(child, id))
            .filter(|result| result.is_ok())
            .map(|result| result.unwrap())
            .next().ok_or(format!("No node with ID {} found", id)),
    }
}

/// Removes merge markers that are inside HTML tags, and replaces the others with  tags.
fn process_merge_markers(html: String) -> String {
    let start_regex = Regex::new(&format!("{}([0-9]+){}", START_MARKER, START_MARKER)).unwrap();
    let end_regex = Regex::new(&format!("{}[0-9]+{}", END_MARKER, END_MARKER)).unwrap();

    let html = remove_merge_markers(html);
    let html = start_regex.replace_all(
        &html, |captures: &Captures| format!("",
                                             captures.at(1).unwrap()));
    end_regex.replace_all(&html, "")
}

fn remove_merge_markers(html: String) -> String {
    // Finds markers where the end, but not the start, is inside a tag.
    let regex1 = Regex::new(&format!(
        r"{}[0-9]+{}([^{}]*?)<([^>]*?){}[0-9]+{}([^>]*?)>",
        START_MARKER, START_MARKER, END_MARKER, END_MARKER, END_MARKER)).unwrap();
    // Finds markers where the start, but not the end, is inside a tag.
    let regex2 = Regex::new(&format!(
        r"<([^>]*?){}[0-9]+{}([^>]*?)>([^{}]*?){}[0-9]+{}",
        START_MARKER, START_MARKER, END_MARKER, END_MARKER, END_MARKER)).unwrap();
    // Finds markers where both the start and end are inside tags.
    let regex3 = Regex::new(&format!(
        r"<([^>]*?){}[0-9]+{}([^>]*?)>([^{}{}]*?)<([^>]*?){}[0-9]+{}([^>]*?)>",
        START_MARKER, START_MARKER, START_MARKER, END_MARKER, END_MARKER, END_MARKER)).unwrap();
    let html = regex1.replace_all(
        &html, |captures: &Captures|
        format!("{}<{}{}>", captures.at(1).unwrap(), captures.at(2).unwrap(),
                captures.at(3).unwrap()));
    let html = regex2.replace_all(
        &html, |captures: &Captures|
        format!("<{}{}>{}", captures.at(1).unwrap(), captures.at(2).unwrap(),
                captures.at(3).unwrap()));
    regex3.replace_all(
        &html, |captures: &Captures|
        format!("<{}{}>{}<{}{}>", captures.at(1).unwrap(), captures.at(2).unwrap(),
                captures.at(3).unwrap(), captures.at(4).unwrap(), captures.at(5).unwrap()))
}

#[cfg(test)]
mod tests {
    use super::{remove_merge_markers, replace_node_with_placeholder};
    use ::START_MARKER;
    use ::END_MARKER;

    fn test_process_merge_markers() {
        let html = format!(
            "{}456{}{}456{}{}123{}t",
            START_MARKER, START_MARKER, END_MARKER, END_MARKER,
            START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected_regex =
            regex!("]*>");
        assert!(expected_regex.is_match(&html));
    }

    #[test]
    fn test_remove_merge_markers_keep() {
        let html = format!("{}456{}{}456{}",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        assert_eq!(html.clone(), remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_keep_one_remove_one() {
        let html = format!(
            "{}234{}text{}234{}{}567{}",
            START_MARKER, START_MARKER, END_MARKER, END_MARKER, START_MARKER, START_MARKER,
            END_MARKER, END_MARKER);
        let expected = format!(
            "{}234{}text{}234{}",
            START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_end_inside_tag() {
        let html = format!("{}123{}",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected = "";
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_start_inside_tag() {
        let html = format!("{}123{}",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected = "";
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_remove_merge_markers_both_inside_tag() {
        let html = format!("text",
                           START_MARKER, START_MARKER, END_MARKER, END_MARKER);
        let expected = "text";
        assert_eq!(expected, remove_merge_markers(html));
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<|endoftext|>


<commit_before>use hyper::Client;

use select::dom::Dom;
use select::predicate::*;

#[derive(Debug)]
pub struct DictDef {
    pub word: String,
    pub def: String,
    pub example: Option<String>,
    pub upvote: i32,
    pub downvote: i32,
    pub contributor: String,
    pub date: String,
    pub id: String,
    pub sounds: Option<Vec<String>>,
    pub similars: Option<Vec<String>>
}

impl DictDef {
}

fn urban_dict_url (word: &str) -> String {
    let mut s = String::new();
    s.push_str("http://www.urbandictionary.com/define.php?term=");
    s.push_str(word);
    s
}

fn json_list_to_strings(json: &str) -> Vec<String> {
    json[1..json.len()-1].split(", ").map(|s| s[1..s.len()-1].to_owned()).collect()
}

fn get_def_from_doc (doc: &Dom) -> Option<DictDef> {
    if let Some(panel) = dom.find(Class("def-panel")).first() {
        if let Some(word) = panel.find(Class("word")).first() {
            let word = word.text();

            let def = panel.find(Class("meaning")).first().unwrap().text();
            let defid = panel.attr("data-defid").unwrap_or("".to_owned());

            let example = panel.find(Class("example")).first().and_then(|e| Some(e.text())).unwrap_or("".to_owned());
            let author = panel.find(Class("author")).first().and_then(|e| Some(e.text())).unwrap_or("".to_owned());
            let date = panel.find(Class("author")).first().and_then(|e| e.next()).and_then(|e| Some(e.text())).unwrap_or("".to_owned());
            let sounds = json_list_to_strings(&(panel.find(Class("play-sound")).first().and_then(|e| e.attr("data-urls")).unwrap_or("[]".to_owned())));

            let similars = doc.find(Name("ul").and(Class("alphabetical"))).find(Name("li")).find(Name("a")).iter().map(|e| e.text());
            Some(DictDef {
                word: word.trim().to_owned(),
                def: def.trim().to_owned(),
                example: example,
                upvote: 0,
                downvote: 0,
                contributor: author,
                date: date,
                id: defid,
                sounds: sounds,
                similars: similars
            })
        } else {
            None
        }
    } else {
        None
    }
}

fn find_from_url(url: &str) -> Option<DictDef> {
    let client = Client::new();

    let mut response = client.get(url).send().unwrap();
    let mut body = String::new();
    resp.read_to_string(&mut body).unwrap();

    let dom = Dom::from_str(body);
    get_def_from_doc(&dom)
}

pub fn find_word_of_the_day() -> Option<DictDef> {
    find_from_url("http://www.urbandictionary.com/")
}

pub fn find_on_urban_dict (word: &str) -> Option<DictDef> {
    let url = urban_dict_url(word);
    find_from_url(&url)
}
<commit_msg>(refactor) refactor with select.rs<commit_after>use std::io::Read;

use hyper::Client;

use select::dom::Dom;
use select::predicate::*;

#[derive(Debug)]
pub struct DictDef {
    pub word: String,
    pub def: String,
    pub example: Option<String>,
    pub upvote: i32,
    pub downvote: i32,
    pub contributor: String,
    pub date: String,
    pub id: String,
    pub sounds: Option<Vec<String>>,
    pub similars: Option<Vec<String>>
}

impl DictDef {
}

fn urban_dict_url (word: &str) -> String {
    let mut s = String::new();
    s.push_str("http://www.urbandictionary.com/define.php?term=");
    s.push_str(word);
    s
}

fn json_list_to_strings(json: &str) -> Vec<String> {
    json[1..json.len()-1].split(", ").map(|s| s[1..s.len()-1].to_owned()).collect()
}

fn get_def_from_doc (doc: &Dom) -> Option<DictDef> {
    if let Some(panel) = doc.find(Class("def-panel")).first() {
        if let Some(word) = panel.find(Class("word")).first() {
            let word = word.text();

            let def = panel.find(Class("meaning")).first().unwrap().text();
            let defid = panel.attr("data-defid").unwrap_or("").to_owned();

            let example = panel.find(Class("example")).first().and_then(|e| Some(e.text()));
            let author = panel.find(Class("author")).first().and_then(|e| Some(e.text())).unwrap_or("".to_owned());
            let date = panel.find(Class("author")).first().and_then(|e| e.next()).and_then(|e| Some(e.text().trim().to_owned())).unwrap_or("".to_owned());

            let sounds = panel.find(Class("play-sound")).first().and_then(|e| e.attr("data-urls").and_then(|s| Some(s.to_owned()))).and_then(|l| Some(json_list_to_strings(&l)));

            let similars = Some(doc.find(Name("ul").and(Class("alphabetical"))).find(Name("li")).find(Name("a")).iter().map(|e| e.text()).collect());

            Some(DictDef {
                word: word.trim().to_owned(),
                def: def.trim().to_owned(),
                example: example,
                upvote: 0,
                downvote: 0,
                contributor: author,
                date: date,
                id: defid,
                sounds: sounds,
                similars: similars
            })
        } else {
            None
        }
    } else {
        None
    }
}

fn find_from_url(url: &str) -> Option<DictDef> {
    let client = Client::new();

    let mut resp = client.get(url).send().unwrap();
    let mut body = String::new();
    {
        resp.read_to_string(&mut body).unwrap();
    }

    let dom = Dom::from_str(&body);
    get_def_from_doc(&dom)
}

pub fn find_word_of_the_day() -> Option<DictDef> {
    find_from_url("http://www.urbandictionary.com/")
}

pub fn find_on_urban_dict (word: &str) -> Option<DictDef> {
    let url = urban_dict_url(word);
    find_from_url(&url)
}
<|endoftext|>


<commit_before>//! Safe-interface corresponding to out-place transform

use super::aligned_vec::*;
use super::enums::*;
use super::plan::*;
use super::r2r::*;

use ndarray::*;
use num_traits::Zero;
use std::marker::PhantomData;

/// Safe-interface corresponding to out-place transform
///
/// FFTW interface modifies an array in `fftw_execute` function
/// which does not takes the array as its arguments.
/// It is not compatible to the programing model of safe Rust.
/// `Pair` interface composes the array and plan to manage
/// mutability in the safe Rust way.
pub struct Pair<A, B, D> {
    pub field: AlignedVec<A>,
    pub coef: AlignedVec<B>,
    logical_size: usize,
    forward: RawPlan,
    backward: RawPlan,
    phantom: PhantomData<D>,
}

impl<A, B, D: Dimension> Pair<A, B, D> {
    pub fn logical_size(&self) -> usize {
        self.logical_size
    }

    /// Execute forward transformation
    pub fn forward(&mut self) {
        unsafe {
            self.forward.execute();
        }
    }

    /// Execute backward transformation
    pub fn backward(&mut self) {
        unsafe {
            self.backward.execute();
        }
    }
}

pub trait ToPair<A, B> {
    type Dim: Dimension;
    fn to_pair(&self) -> Pair<A, B, Self::Dim>;
}

#[derive(Debug, Clone, Copy)]
pub struct R2R1D {
    n: usize,
    kind: R2R_KIND,
    flag: FLAG,
}

impl<T: R2R + AlignedAllocable + Zero> ToPair<T, T> for R2R1D {
    type Dim = Ix1;
    fn to_pair(&self) -> Pair<T, T, Ix1> {
        let mut field = AlignedVec::new(self.n);
        let mut coef = AlignedVec::new(self.n);
        let forward = unsafe { T::r2r_1d(self.n, &mut field, &mut coef, forward(self.kind), self.flag) };
        let backward = unsafe {
            T::r2r_1d(
                self.n,
                &mut coef,
                &mut field,
                backward(self.kind),
                self.flag,
            )
        };
        Pair {
            field: field,
            coef: coef,
            logical_size: logical_size(self.n, self.kind),
            forward: forward,
            backward: backward,
            phantom: PhantomData,
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct C2C1D {
    n: usize,
    sign: SIGN,
    flag: FLAG,
}

impl<T: C2C + AlignedAllocable + Zero> ToPair<T, T> for C2C1D {
    type Dim = Ix1;
    fn to_pair(&self) -> Pair<T, T, Ix1> {
        let mut field = AlignedVec::new(self.n);
        let mut coef = AlignedVec::new(self.n);
        let forward = unsafe { T::c2c_1d(self.n, &mut field, &mut coef, self.sign, self.flag) };
        let backward = unsafe { T::c2c_1d(self.n, &mut coef, &mut field, -self.sign, self.flag) };
        Pair {
            field: field,
            coef: coef,
            logical_size: self.n,
            forward: forward,
            backward: backward,
            phantom: PhantomData,
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct R2C1D {
    n: usize,
    flag: FLAG,
}

impl<R, C> ToPair<R, C> for R2C1D
where
    (C, R): C2R<Real = R, Complex = C>,
    R: AlignedAllocable + Zero,
    C: AlignedAllocable + Zero,
{
    type Dim = Ix1;
    fn to_pair(&self) -> Pair<R, C, Ix1> {
        let mut field = AlignedVec::<R>::new(self.n);
        let mut coef = AlignedVec::<C>::new(self.n / 2 + 1);
        let forward = unsafe { <(C, R) as C2R>::r2c_1d(self.n, &mut field, &mut coef, self.flag) };
        let backward = unsafe { <(C, R) as C2R>::c2r_1d(self.n, &mut coef, &mut field, self.flag) };
        Pair {
            field: field,
            coef: coef,
            logical_size: self.n,
            forward: forward,
            backward: backward,
            phantom: PhantomData,
        }
    }
}
<commit_msg>Add utility functions<commit_after>//! Safe-interface corresponding to out-place transform

use super::aligned_vec::*;
use super::enums::*;
use super::plan::*;
use super::r2r::*;

use ndarray::*;
use num_traits::Zero;
use std::marker::PhantomData;

/// Safe-interface corresponding to out-place transform
///
/// FFTW interface modifies an array in `fftw_execute` function
/// which does not takes the array as its arguments.
/// It is not compatible to the programing model of safe Rust.
/// `Pair` interface composes the array and plan to manage
/// mutability in the safe Rust way.
pub struct Pair<A, B, D> {
    pub field: AlignedVec<A>,
    pub coef: AlignedVec<B>,
    logical_size: usize,
    forward: RawPlan,
    backward: RawPlan,
    phantom: PhantomData<D>,
}

impl<A, B, D: Dimension> Pair<A, B, D> {
    pub fn logical_size(&self) -> usize {
        self.logical_size
    }

    /// Execute forward transformation
    pub fn forward(&mut self) {
        unsafe {
            self.forward.execute();
        }
    }

    /// Execute backward transformation
    pub fn backward(&mut self) {
        unsafe {
            self.backward.execute();
        }
    }
}

/// Create a `Pair` from a setting struct e.g. `R2C1D`.
pub trait ToPair<A, B> {
    type Dim: Dimension;
    /// Generate `Pair` from a setting struct
    fn to_pair(&self) -> Pair<A, B, Self::Dim>;
}

#[derive(Debug, Clone, Copy)]
pub struct R2R1D {
    n: usize,
    kind: R2R_KIND,
    flag: FLAG,
}

impl<T: R2R + AlignedAllocable + Zero> ToPair<T, T> for R2R1D {
    type Dim = Ix1;
    fn to_pair(&self) -> Pair<T, T, Ix1> {
        let mut field = AlignedVec::new(self.n);
        let mut coef = AlignedVec::new(self.n);
        let forward = unsafe { T::r2r_1d(self.n, &mut field, &mut coef, forward(self.kind), self.flag) };
        let backward = unsafe {
            T::r2r_1d(
                self.n,
                &mut coef,
                &mut field,
                backward(self.kind),
                self.flag,
            )
        };
        Pair {
            field: field,
            coef: coef,
            logical_size: logical_size(self.n, self.kind),
            forward: forward,
            backward: backward,
            phantom: PhantomData,
        }
    }
}

/// Setting for 1-dimensional C2C transform
#[derive(Debug, Clone, Copy)]
pub struct C2C1D {
    n: usize,
    sign: SIGN,
    flag: FLAG,
}

/// Utility function to generage 1-dimensional C2C setting
pub fn c2c_1d(n: usize) -> C2C1D {
    C2C1D {
        n,
        sign: SIGN::FFTW_FORWARD,
        flag: FLAG::FFTW_MEASURE,
    }
}

impl<T: C2C + AlignedAllocable + Zero> ToPair<T, T> for C2C1D {
    type Dim = Ix1;
    fn to_pair(&self) -> Pair<T, T, Ix1> {
        let mut field = AlignedVec::new(self.n);
        let mut coef = AlignedVec::new(self.n);
        let forward = unsafe { T::c2c_1d(self.n, &mut field, &mut coef, self.sign, self.flag) };
        let backward = unsafe { T::c2c_1d(self.n, &mut coef, &mut field, -self.sign, self.flag) };
        Pair {
            field: field,
            coef: coef,
            logical_size: self.n,
            forward: forward,
            backward: backward,
            phantom: PhantomData,
        }
    }
}

/// Setting for 1-dimensional R2C transform
#[derive(Debug, Clone, Copy)]
pub struct R2C1D {
    n: usize,
    flag: FLAG,
}

/// Utility function to generage 1-dimensional R2C setting
pub fn r2c_1d(n: usize) -> R2C1D {
    R2C1D {
        n,
        flag: FLAG::FFTW_MEASURE,
    }
}

impl<R, C> ToPair<R, C> for R2C1D
where
    (C, R): C2R<Real = R, Complex = C>,
    R: AlignedAllocable + Zero,
    C: AlignedAllocable + Zero,
{
    type Dim = Ix1;
    fn to_pair(&self) -> Pair<R, C, Ix1> {
        let mut field = AlignedVec::<R>::new(self.n);
        let mut coef = AlignedVec::<C>::new(self.n / 2 + 1);
        let forward = unsafe { <(C, R) as C2R>::r2c_1d(self.n, &mut field, &mut coef, self.flag) };
        let backward = unsafe { <(C, R) as C2R>::c2r_1d(self.n, &mut coef, &mut field, self.flag) };
        Pair {
            field: field,
            coef: coef,
            logical_size: self.n,
            forward: forward,
            backward: backward,
            phantom: PhantomData,
        }
    }
}
<|endoftext|>


<commit_before>use native;
use std::comm::channel;

struct WorkUnit {
    arg: int,
    callback: Sender<int>,
}

enum MessageToWorker {
    Work(WorkUnit),
    Halt,
}

enum MessageToDispatcher {
    Dispatch(WorkUnit),
    HaltOne,
}

struct WorkQueue {
    dispatcher: Sender<MessageToDispatcher>,
    worker_count: int,
}

impl WorkQueue {
    pub fn create(worker_count: int) -> WorkQueue {
        let (want_work, idle_worker) = channel::<Sender<MessageToWorker>>();
        let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher>();

        for _ in range(0, worker_count) {
            // worker
            let want_work_copy = want_work.clone();
            native::task::spawn(proc() {
                let want_work = want_work_copy;
                loop {
                    let (reply_with_work, get_work_unit) = channel::<MessageToWorker>();
                    want_work.send(reply_with_work);
                    let work_unit = match get_work_unit.recv() {
                        Work(wu) => wu,
                        Halt     => return
                    };
                    let rv = work_unit.arg * 2;
                    work_unit.callback.send(rv);
                }
            });
        }

        // dispatcher
        native::task::spawn(proc() {
            let mut worker_count = worker_count;
            let inbox = dispatcher_inbox;
            let idle_worker = idle_worker;
            loop {
                match inbox.recv() {
                    Dispatch(work_item) => {
                        idle_worker.recv().send(Work(work_item));
                    },
                    HaltOne => {
                        worker_count -= 1;
                        idle_worker.recv().send(Halt);
                        if worker_count == 0 {
                            return;
                        }
                    },
                };
            }
        });

        return WorkQueue{dispatcher: dispatcher, worker_count: worker_count};
    }

    pub fn execute(&self, arg: int) -> Receiver<int> {
        let (reply_with_rv, wait_for_rv) = channel::<int>();
        self.dispatcher.send(Dispatch(WorkUnit{arg: arg, callback: reply_with_rv}));
        return wait_for_rv;
    }
}

impl Drop for WorkQueue {
    fn drop(&mut self) {
        for _ in range(0, self.worker_count) {
            self.dispatcher.send(HaltOne);
        }
    }
}

#[test]
fn test_queue() {
    let queue = WorkQueue::create(3);
    let mut promise_list: ~[Receiver<int>] = ~[];
    for c in range(0, 10) {
        let rv = queue.execute(c);
        promise_list.push(rv);
    }
    let return_list = promise_list.map(|promise| promise.recv());
    assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
}
<commit_msg>dispatcher knows how many tasks to halt<commit_after>use native;
use std::comm::channel;

struct WorkUnit {
    arg: int,
    callback: Sender<int>,
}

enum MessageToWorker {
    Work(WorkUnit),
    Halt,
}

enum MessageToDispatcher {
    Dispatch(WorkUnit),
    HaltAll,
}

struct WorkQueue {
    dispatcher: Sender<MessageToDispatcher>,
}

impl WorkQueue {
    pub fn create(worker_count: int) -> WorkQueue {
        let (want_work, idle_worker) = channel::<Sender<MessageToWorker>>();
        let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher>();

        for _ in range(0, worker_count) {
            // worker
            let want_work_copy = want_work.clone();
            native::task::spawn(proc() {
                let want_work = want_work_copy;
                loop {
                    let (reply_with_work, get_work_unit) = channel::<MessageToWorker>();
                    want_work.send(reply_with_work);
                    let work_unit = match get_work_unit.recv() {
                        Work(wu) => wu,
                        Halt     => return
                    };
                    let rv = work_unit.arg * 2;
                    work_unit.callback.send(rv);
                }
            });
        }

        // dispatcher
        native::task::spawn(proc() {
            let mut worker_count = worker_count;
            let inbox = dispatcher_inbox;
            let idle_worker = idle_worker;
            loop {
                match inbox.recv() {
                    Dispatch(work_item) => {
                        idle_worker.recv().send(Work(work_item));
                    },
                    HaltAll => {
                        while worker_count > 0 {
                            idle_worker.recv().send(Halt);
                            worker_count -= 1;
                        }
                        return;
                    },
                };
            }
        });

        return WorkQueue{dispatcher: dispatcher};
    }

    pub fn execute(&self, arg: int) -> Receiver<int> {
        let (reply_with_rv, wait_for_rv) = channel::<int>();
        self.dispatcher.send(Dispatch(WorkUnit{arg: arg, callback: reply_with_rv}));
        return wait_for_rv;
    }
}

impl Drop for WorkQueue {
    fn drop(&mut self) {
        self.dispatcher.send(HaltAll);
    }
}

#[test]
fn test_queue() {
    let queue = WorkQueue::create(3);
    let mut promise_list: ~[Receiver<int>] = ~[];
    for c in range(0, 10) {
        let rv = queue.execute(c);
        promise_list.push(rv);
    }
    let return_list = promise_list.map(|promise| promise.recv());
    assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
}
<|endoftext|>