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

#![deny(warnings)]

extern crate filetime;

use std::fs::File;
use std::io;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::{fs, env};

use filetime::FileTime;

/// A helper macro to `unwrap` a result except also print out details like:
///
/// * The file/line of the panic
/// * The expression that failed
/// * The error itself
///
/// This is currently used judiciously throughout the build system rather than
/// using a `Result` with `try!`, but this may change one day...
#[macro_export]
macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {}", stringify!($e), e),
    })
}

pub fn run(cmd: &mut Command) {
    println!("running: {:?}", cmd);
    run_silent(cmd);
}

pub fn run_silent(cmd: &mut Command) {
    if !try_run_silent(cmd) {
        std::process::exit(1);
    }
}

pub fn try_run_silent(cmd: &mut Command) -> bool {
    let status = match cmd.status() {
        Ok(status) => status,
        Err(e) => fail(&format!("failed to execute command: {:?}\nerror: {}",
                                cmd, e)),
    };
    if !status.success() {
        println!("\n\ncommand did not execute successfully: {:?}\n\
                  expected success, got: {}\n\n",
                 cmd,
                 status);
    }
    status.success()
}

pub fn run_suppressed(cmd: &mut Command) {
    if !try_run_suppressed(cmd) {
        std::process::exit(1);
    }
}

pub fn try_run_suppressed(cmd: &mut Command) -> bool {
    let output = match cmd.output() {
        Ok(status) => status,
        Err(e) => fail(&format!("failed to execute command: {:?}\nerror: {}",
                                cmd, e)),
    };
    if !output.status.success() {
        println!("\n\ncommand did not execute successfully: {:?}\n\
                  expected success, got: {}\n\n\
                  stdout ----\n{}\n\
                  stderr ----\n{}\n\n",
                 cmd,
                 output.status,
                 String::from_utf8_lossy(&output.stdout),
                 String::from_utf8_lossy(&output.stderr));
    }
    output.status.success()
}

pub fn gnu_target(target: &str) -> String {
    match target {
        "i686-pc-windows-msvc" => "i686-pc-win32".to_string(),
        "x86_64-pc-windows-msvc" => "x86_64-pc-win32".to_string(),
        "i686-pc-windows-gnu" => "i686-w64-mingw32".to_string(),
        "x86_64-pc-windows-gnu" => "x86_64-w64-mingw32".to_string(),
        s => s.to_string(),
    }
}

pub fn cc2ar(cc: &Path, target: &str) -> Option<PathBuf> {
    if target.contains("msvc") {
        None
    } else if target.contains("musl") {
        Some(PathBuf::from("ar"))
    } else if target.contains("openbsd") {
        Some(PathBuf::from("ar"))
    } else {
        let parent = cc.parent().unwrap();
        let file = cc.file_name().unwrap().to_str().unwrap();
        for suffix in &["gcc", "cc", "clang"] {
            if let Some(idx) = file.rfind(suffix) {
                let mut file = file[..idx].to_owned();
                file.push_str("ar");
                return Some(parent.join(&file));
            }
        }
        Some(parent.join(file))
    }
}

pub fn make(host: &str) -> PathBuf {
    if host.contains("bitrig") || host.contains("dragonfly") ||
        host.contains("freebsd") || host.contains("netbsd") ||
        host.contains("openbsd") {
        PathBuf::from("gmake")
    } else {
        PathBuf::from("make")
    }
}

pub fn output(cmd: &mut Command) -> String {
    let output = match cmd.stderr(Stdio::inherit()).output() {
        Ok(status) => status,
        Err(e) => fail(&format!("failed to execute command: {:?}\nerror: {}",
                                cmd, e)),
    };
    if !output.status.success() {
        panic!("command did not execute successfully: {:?}\n\
                expected success, got: {}",
               cmd,
               output.status);
    }
    String::from_utf8(output.stdout).unwrap()
}

pub fn rerun_if_changed_anything_in_dir(dir: &Path) {
    let mut stack = dir.read_dir().unwrap()
                       .map(|e| e.unwrap())
                       .filter(|e| &*e.file_name() != ".git")
                       .collect::<Vec<_>>();
    while let Some(entry) = stack.pop() {
        let path = entry.path();
        if entry.file_type().unwrap().is_dir() {
            stack.extend(path.read_dir().unwrap().map(|e| e.unwrap()));
        } else {
            println!("cargo:rerun-if-changed={}", path.display());
        }
    }
}

/// Returns the last-modified time for `path`, or zero if it doesn't exist.
pub fn mtime(path: &Path) -> FileTime {
    fs::metadata(path).map(|f| {
        FileTime::from_last_modification_time(&f)
    }).unwrap_or(FileTime::zero())
}

/// Returns whether `dst` is up to date given that the file or files in `src`
/// are used to generate it.
///
/// Uses last-modified time checks to verify this.
pub fn up_to_date(src: &Path, dst: &Path) -> bool {
    let threshold = mtime(dst);
    let meta = match fs::metadata(src) {
        Ok(meta) => meta,
        Err(e) => panic!("source {:?} failed to get metadata: {}", src, e),
    };
    if meta.is_dir() {
        dir_up_to_date(src, &threshold)
    } else {
        FileTime::from_last_modification_time(&meta) <= threshold
    }
}

#[must_use]
pub struct NativeLibBoilerplate {
    pub src_dir: PathBuf,
    pub out_dir: PathBuf,
}

impl Drop for NativeLibBoilerplate {
    fn drop(&mut self) {
        t!(File::create(self.out_dir.join("rustbuild.timestamp")));
    }
}

// Perform standard preparations for native libraries that are build only once for all stages.
// Emit rerun-if-changed and linking attributes for Cargo, check if any source files are
// updated, calculate paths used later in actual build with CMake/make or C/C++ compiler.
// If Err is returned, then everything is up-to-date and further build actions can be skipped.
// Timestamps are created automatically when the result of `native_lib_boilerplate` goes out
// of scope, so all the build actions should be completed until then.
pub fn native_lib_boilerplate(src_name: &str,
                              out_name: &str,
                              link_name: &str,
                              search_subdir: &str)
                              -> Result<NativeLibBoilerplate, ()> {
    let current_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    let src_dir = current_dir.join("..").join(src_name);
    rerun_if_changed_anything_in_dir(&src_dir);

    let out_dir = env::var_os("RUSTBUILD_NATIVE_DIR").unwrap_or(env::var_os("OUT_DIR").unwrap());
    let out_dir = PathBuf::from(out_dir).join(out_name);
    t!(create_dir_racy(&out_dir));
    if link_name.contains('=') {
        println!("cargo:rustc-link-lib={}", link_name);
    } else {
        println!("cargo:rustc-link-lib=static={}", link_name);
    }
    println!("cargo:rustc-link-search=native={}", out_dir.join(search_subdir).display());

    let timestamp = out_dir.join("rustbuild.timestamp");
    if !up_to_date(Path::new("build.rs"), &timestamp) || !up_to_date(&src_dir, &timestamp) {
        Ok(NativeLibBoilerplate { src_dir: src_dir, out_dir: out_dir })
    } else {
        Err(())
    }
}

pub fn sanitizer_lib_boilerplate(sanitizer_name: &str) -> Result<NativeLibBoilerplate, ()> {
    let (link_name, search_path) = match &*env::var("TARGET").unwrap() {
        "x86_64-unknown-linux-gnu" => (
            format!("clang_rt.{}-x86_64", sanitizer_name),
            "build/lib/linux",
        ),
        "x86_64-apple-darwin" => (
            format!("dylib=clang_rt.{}_osx_dynamic", sanitizer_name),
            "build/lib/darwin",
        ),
        _ => return Err(()),
    };
    native_lib_boilerplate("libcompiler_builtins/compiler-rt",
                           sanitizer_name,
                           &link_name,
                           search_path)
}

fn dir_up_to_date(src: &Path, threshold: &FileTime) -> bool {
    t!(fs::read_dir(src)).map(|e| t!(e)).all(|e| {
        let meta = t!(e.metadata());
        if meta.is_dir() {
            dir_up_to_date(&e.path(), threshold)
        } else {
            FileTime::from_last_modification_time(&meta) < *threshold
        }
    })
}

fn fail(s: &str) -> ! {
    println!("\n\n{}\n\n", s);
    std::process::exit(1);
}

fn create_dir_racy(path: &Path) -> io::Result<()> {
    match fs::create_dir(path) {
        Ok(()) => return Ok(()),
        Err(ref e) if e.kind() == io::ErrorKind::AlreadyExists => return Ok(()),
        Err(ref e) if e.kind() == io::ErrorKind::NotFound => {}
        Err(e) => return Err(e),
    }
    match path.parent() {
        Some(p) => try!(create_dir_racy(p)),
        None => return Err(io::Error::new(io::ErrorKind::Other, "failed to create whole tree")),
    }
    match fs::create_dir(path) {
        Ok(()) => Ok(()),
        Err(ref e) if e.kind() == io::ErrorKind::AlreadyExists => Ok(()),
        Err(e) => Err(e),
    }
}
<commit_msg>rustbuild: Replace create_dir_racy with create_dir_all<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![deny(warnings)]

extern crate filetime;

use std::fs::File;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::{fs, env};

use filetime::FileTime;

/// A helper macro to `unwrap` a result except also print out details like:
///
/// * The file/line of the panic
/// * The expression that failed
/// * The error itself
///
/// This is currently used judiciously throughout the build system rather than
/// using a `Result` with `try!`, but this may change one day...
#[macro_export]
macro_rules! t {
    ($e:expr) => (match $e {
        Ok(e) => e,
        Err(e) => panic!("{} failed with {}", stringify!($e), e),
    })
}

pub fn run(cmd: &mut Command) {
    println!("running: {:?}", cmd);
    run_silent(cmd);
}

pub fn run_silent(cmd: &mut Command) {
    if !try_run_silent(cmd) {
        std::process::exit(1);
    }
}

pub fn try_run_silent(cmd: &mut Command) -> bool {
    let status = match cmd.status() {
        Ok(status) => status,
        Err(e) => fail(&format!("failed to execute command: {:?}\nerror: {}",
                                cmd, e)),
    };
    if !status.success() {
        println!("\n\ncommand did not execute successfully: {:?}\n\
                  expected success, got: {}\n\n",
                 cmd,
                 status);
    }
    status.success()
}

pub fn run_suppressed(cmd: &mut Command) {
    if !try_run_suppressed(cmd) {
        std::process::exit(1);
    }
}

pub fn try_run_suppressed(cmd: &mut Command) -> bool {
    let output = match cmd.output() {
        Ok(status) => status,
        Err(e) => fail(&format!("failed to execute command: {:?}\nerror: {}",
                                cmd, e)),
    };
    if !output.status.success() {
        println!("\n\ncommand did not execute successfully: {:?}\n\
                  expected success, got: {}\n\n\
                  stdout ----\n{}\n\
                  stderr ----\n{}\n\n",
                 cmd,
                 output.status,
                 String::from_utf8_lossy(&output.stdout),
                 String::from_utf8_lossy(&output.stderr));
    }
    output.status.success()
}

pub fn gnu_target(target: &str) -> String {
    match target {
        "i686-pc-windows-msvc" => "i686-pc-win32".to_string(),
        "x86_64-pc-windows-msvc" => "x86_64-pc-win32".to_string(),
        "i686-pc-windows-gnu" => "i686-w64-mingw32".to_string(),
        "x86_64-pc-windows-gnu" => "x86_64-w64-mingw32".to_string(),
        s => s.to_string(),
    }
}

pub fn cc2ar(cc: &Path, target: &str) -> Option<PathBuf> {
    if target.contains("msvc") {
        None
    } else if target.contains("musl") {
        Some(PathBuf::from("ar"))
    } else if target.contains("openbsd") {
        Some(PathBuf::from("ar"))
    } else {
        let parent = cc.parent().unwrap();
        let file = cc.file_name().unwrap().to_str().unwrap();
        for suffix in &["gcc", "cc", "clang"] {
            if let Some(idx) = file.rfind(suffix) {
                let mut file = file[..idx].to_owned();
                file.push_str("ar");
                return Some(parent.join(&file));
            }
        }
        Some(parent.join(file))
    }
}

pub fn make(host: &str) -> PathBuf {
    if host.contains("bitrig") || host.contains("dragonfly") ||
        host.contains("freebsd") || host.contains("netbsd") ||
        host.contains("openbsd") {
        PathBuf::from("gmake")
    } else {
        PathBuf::from("make")
    }
}

pub fn output(cmd: &mut Command) -> String {
    let output = match cmd.stderr(Stdio::inherit()).output() {
        Ok(status) => status,
        Err(e) => fail(&format!("failed to execute command: {:?}\nerror: {}",
                                cmd, e)),
    };
    if !output.status.success() {
        panic!("command did not execute successfully: {:?}\n\
                expected success, got: {}",
               cmd,
               output.status);
    }
    String::from_utf8(output.stdout).unwrap()
}

pub fn rerun_if_changed_anything_in_dir(dir: &Path) {
    let mut stack = dir.read_dir().unwrap()
                       .map(|e| e.unwrap())
                       .filter(|e| &*e.file_name() != ".git")
                       .collect::<Vec<_>>();
    while let Some(entry) = stack.pop() {
        let path = entry.path();
        if entry.file_type().unwrap().is_dir() {
            stack.extend(path.read_dir().unwrap().map(|e| e.unwrap()));
        } else {
            println!("cargo:rerun-if-changed={}", path.display());
        }
    }
}

/// Returns the last-modified time for `path`, or zero if it doesn't exist.
pub fn mtime(path: &Path) -> FileTime {
    fs::metadata(path).map(|f| {
        FileTime::from_last_modification_time(&f)
    }).unwrap_or(FileTime::zero())
}

/// Returns whether `dst` is up to date given that the file or files in `src`
/// are used to generate it.
///
/// Uses last-modified time checks to verify this.
pub fn up_to_date(src: &Path, dst: &Path) -> bool {
    let threshold = mtime(dst);
    let meta = match fs::metadata(src) {
        Ok(meta) => meta,
        Err(e) => panic!("source {:?} failed to get metadata: {}", src, e),
    };
    if meta.is_dir() {
        dir_up_to_date(src, &threshold)
    } else {
        FileTime::from_last_modification_time(&meta) <= threshold
    }
}

#[must_use]
pub struct NativeLibBoilerplate {
    pub src_dir: PathBuf,
    pub out_dir: PathBuf,
}

impl Drop for NativeLibBoilerplate {
    fn drop(&mut self) {
        t!(File::create(self.out_dir.join("rustbuild.timestamp")));
    }
}

// Perform standard preparations for native libraries that are build only once for all stages.
// Emit rerun-if-changed and linking attributes for Cargo, check if any source files are
// updated, calculate paths used later in actual build with CMake/make or C/C++ compiler.
// If Err is returned, then everything is up-to-date and further build actions can be skipped.
// Timestamps are created automatically when the result of `native_lib_boilerplate` goes out
// of scope, so all the build actions should be completed until then.
pub fn native_lib_boilerplate(src_name: &str,
                              out_name: &str,
                              link_name: &str,
                              search_subdir: &str)
                              -> Result<NativeLibBoilerplate, ()> {
    let current_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    let src_dir = current_dir.join("..").join(src_name);
    rerun_if_changed_anything_in_dir(&src_dir);

    let out_dir = env::var_os("RUSTBUILD_NATIVE_DIR").unwrap_or(env::var_os("OUT_DIR").unwrap());
    let out_dir = PathBuf::from(out_dir).join(out_name);
    t!(fs::create_dir_all(&out_dir));
    if link_name.contains('=') {
        println!("cargo:rustc-link-lib={}", link_name);
    } else {
        println!("cargo:rustc-link-lib=static={}", link_name);
    }
    println!("cargo:rustc-link-search=native={}", out_dir.join(search_subdir).display());

    let timestamp = out_dir.join("rustbuild.timestamp");
    if !up_to_date(Path::new("build.rs"), &timestamp) || !up_to_date(&src_dir, &timestamp) {
        Ok(NativeLibBoilerplate { src_dir: src_dir, out_dir: out_dir })
    } else {
        Err(())
    }
}

pub fn sanitizer_lib_boilerplate(sanitizer_name: &str) -> Result<NativeLibBoilerplate, ()> {
    let (link_name, search_path) = match &*env::var("TARGET").unwrap() {
        "x86_64-unknown-linux-gnu" => (
            format!("clang_rt.{}-x86_64", sanitizer_name),
            "build/lib/linux",
        ),
        "x86_64-apple-darwin" => (
            format!("dylib=clang_rt.{}_osx_dynamic", sanitizer_name),
            "build/lib/darwin",
        ),
        _ => return Err(()),
    };
    native_lib_boilerplate("libcompiler_builtins/compiler-rt",
                           sanitizer_name,
                           &link_name,
                           search_path)
}

fn dir_up_to_date(src: &Path, threshold: &FileTime) -> bool {
    t!(fs::read_dir(src)).map(|e| t!(e)).all(|e| {
        let meta = t!(e.metadata());
        if meta.is_dir() {
            dir_up_to_date(&e.path(), threshold)
        } else {
            FileTime::from_last_modification_time(&meta) < *threshold
        }
    })
}

fn fail(s: &str) -> ! {
    println!("\n\n{}\n\n", s);
    std::process::exit(1);
}
<|endoftext|>


<commit_before>use crate::core::shell::Verbosity;
use crate::core::{GitReference, Workspace};
use crate::ops;
use crate::sources::path::PathSource;
use crate::sources::CRATES_IO_REGISTRY;
use crate::util::{CargoResult, Config};
use anyhow::{bail, Context as _};
use cargo_util::{paths, Sha256};
use serde::Serialize;
use std::collections::HashSet;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use toml_edit::easy as toml;

pub struct VendorOptions<'a> {
    pub no_delete: bool,
    pub versioned_dirs: bool,
    pub destination: &'a Path,
    pub extra: Vec<PathBuf>,
}

pub fn vendor(ws: &Workspace<'_>, opts: &VendorOptions<'_>) -> CargoResult<()> {
    let config = ws.config();
    let mut extra_workspaces = Vec::new();
    for extra in opts.extra.iter() {
        let extra = config.cwd().join(extra);
        let ws = Workspace::new(&extra, config)?;
        extra_workspaces.push(ws);
    }
    let workspaces = extra_workspaces.iter().chain(Some(ws)).collect::<Vec<_>>();
    let vendor_config = sync(config, &workspaces, opts).with_context(|| "failed to sync")?;

    if config.shell().verbosity() != Verbosity::Quiet {
        if vendor_config.source.is_empty() {
            crate::drop_eprintln!(config, "There is no dependency to vendor in this project.");
        } else {
            crate::drop_eprint!(
                config,
                "To use vendored sources, add this to your .cargo/config.toml for this project:\n\n"
            );
            crate::drop_print!(
                config,
                "{}",
                &toml::to_string_pretty(&vendor_config).unwrap()
            );
        }
    }

    Ok(())
}

#[derive(Serialize)]
struct VendorConfig {
    source: BTreeMap<String, VendorSource>,
}

#[derive(Serialize)]
#[serde(rename_all = "lowercase", untagged)]
enum VendorSource {
    Directory {
        directory: String,
    },
    Registry {
        registry: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
    Git {
        git: String,
        branch: Option<String>,
        tag: Option<String>,
        rev: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
}

fn sync(
    config: &Config,
    workspaces: &[&Workspace<'_>],
    opts: &VendorOptions<'_>,
) -> CargoResult<VendorConfig> {
    let canonical_destination = opts.destination.canonicalize();
    let canonical_destination = canonical_destination.as_deref().unwrap_or(opts.destination);
    let dest_dir_already_exists = canonical_destination.exists();

    paths::create_dir_all(&canonical_destination)?;
    let mut to_remove = HashSet::new();
    if !opts.no_delete {
        for entry in canonical_destination.read_dir()? {
            let entry = entry?;
            if !entry
                .file_name()
                .to_str()
                .map_or(false, |s| s.starts_with('.'))
            {
                to_remove.insert(entry.path());
            }
        }
    }

    // First up attempt to work around rust-lang/cargo#5956. Apparently build
    // artifacts sprout up in Cargo's global cache for whatever reason, although
    // it's unsure what tool is causing these issues at this time. For now we
    // apply a heavy-hammer approach which is to delete Cargo's unpacked version
    // of each crate to start off with. After we do this we'll re-resolve and
    // redownload again, which should trigger Cargo to re-extract all the
    // crates.
    //
    // Note that errors are largely ignored here as this is a best-effort
    // attempt. If anything fails here we basically just move on to the next
    // crate to work with.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).with_context(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .with_context(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // Don't delete actual source code!
            if pkg.source_id().is_path() {
                if let Ok(path) = pkg.source_id().url().to_file_path() {
                    if let Ok(path) = path.canonicalize() {
                        to_remove.remove(&path);
                    }
                }
                continue;
            }
            if pkg.source_id().is_git() {
                continue;
            }
            if let Ok(pkg) = packages.get_one(pkg) {
                drop(fs::remove_dir_all(pkg.manifest_path().parent().unwrap()));
            }
        }
    }

    let mut checksums = HashMap::new();
    let mut ids = BTreeMap::new();

    // Next up let's actually download all crates and start storing internal
    // tables about them.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).with_context(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .with_context(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // No need to vendor path crates since they're already in the
            // repository
            if pkg.source_id().is_path() {
                continue;
            }
            ids.insert(
                pkg,
                packages
                    .get_one(pkg)
                    .with_context(|| "failed to fetch package")?
                    .clone(),
            );

            checksums.insert(pkg, resolve.checksums().get(&pkg).cloned());
        }
    }

    let mut versions = HashMap::new();
    for id in ids.keys() {
        let map = versions.entry(id.name()).or_insert_with(BTreeMap::default);
        if let Some(prev) = map.get(&id.version()) {
            bail!(
                "found duplicate version of package `{} v{}` \
                 vendored from two sources:\n\
                 \n\
                 \tsource 1: {}\n\
                 \tsource 2: {}",
                id.name(),
                id.version(),
                prev,
                id.source_id()
            );
        }
        map.insert(id.version(), id.source_id());
    }

    let mut sources = BTreeSet::new();
    let mut tmp_buf = [0; 64 * 1024];
    for (id, pkg) in ids.iter() {
        // Next up, copy it to the vendor directory
        let src = pkg
            .manifest_path()
            .parent()
            .expect("manifest_path should point to a file");
        let max_version = *versions[&id.name()].iter().rev().next().unwrap().0;
        let dir_has_version_suffix = opts.versioned_dirs || id.version() != max_version;
        let dst_name = if dir_has_version_suffix {
            // Eg vendor/futures-0.1.13
            format!("{}-{}", id.name(), id.version())
        } else {
            // Eg vendor/futures
            id.name().to_string()
        };

        sources.insert(id.source_id());
        let dst = canonical_destination.join(&dst_name);
        to_remove.remove(&dst);
        let cksum = dst.join(".cargo-checksum.json");
        if dir_has_version_suffix && cksum.exists() {
            // Always re-copy directory without version suffix in case the version changed
            continue;
        }

        config.shell().status(
            "Vendoring",
            &format!("{} ({}) to {}", id, src.to_string_lossy(), dst.display()),
        )?;

        let _ = fs::remove_dir_all(&dst);
        let pathsource = PathSource::new(src, id.source_id(), config);
        let paths = pathsource.list_files(pkg)?;
        let mut map = BTreeMap::new();
        cp_sources(src, &paths, &dst, &mut map, &mut tmp_buf)
            .with_context(|| format!("failed to copy over vendored sources for: {}", id))?;

        // Finally, emit the metadata about this package
        let json = serde_json::json!({
            "package": checksums.get(id),
            "files": map,
        });

        paths::write(&cksum, json.to_string())?;
    }

    for path in to_remove {
        if path.is_dir() {
            paths::remove_dir_all(&path)?;
        } else {
            paths::remove_file(&path)?;
        }
    }

    // add our vendored source
    let mut config = BTreeMap::new();

    let merged_source_name = "vendored-sources";

    // replace original sources with vendor
    for source_id in sources {
        let name = if source_id.is_default_registry() {
            CRATES_IO_REGISTRY.to_string()
        } else {
            source_id.url().to_string()
        };

        let source = if source_id.is_default_registry() {
            VendorSource::Registry {
                registry: None,
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_remote_registry() {
            let registry = source_id.url().to_string();
            VendorSource::Registry {
                registry: Some(registry),
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_git() {
            let mut branch = None;
            let mut tag = None;
            let mut rev = None;
            if let Some(reference) = source_id.git_reference() {
                match *reference {
                    GitReference::Branch(ref b) => branch = Some(b.clone()),
                    GitReference::Tag(ref t) => tag = Some(t.clone()),
                    GitReference::Rev(ref r) => rev = Some(r.clone()),
                    GitReference::DefaultBranch => {}
                }
            }
            VendorSource::Git {
                git: source_id.url().to_string(),
                branch,
                tag,
                rev,
                replace_with: merged_source_name.to_string(),
            }
        } else {
            panic!("Invalid source ID: {}", source_id)
        };
        config.insert(name, source);
    }

    if !config.is_empty() {
        config.insert(
            merged_source_name.to_string(),
            VendorSource::Directory {
                directory: opts.destination.to_string_lossy().replace("\\", "/"),
            },
        );
    } else if !dest_dir_already_exists {
        // Nothing to vendor. Remove the destination dir we've just created.
        paths::remove_dir(canonical_destination)?;
    }

    Ok(VendorConfig { source: config })
}

fn cp_sources(
    src: &Path,
    paths: &[PathBuf],
    dst: &Path,
    cksums: &mut BTreeMap<String, String>,
    tmp_buf: &mut [u8],
) -> CargoResult<()> {
    for p in paths {
        let relative = p.strip_prefix(&src).unwrap();

        match relative.to_str() {
            // Skip git config files as they're not relevant to builds most of
            // the time and if we respect them (e.g.  in git) then it'll
            // probably mess with the checksums when a vendor dir is checked
            // into someone else's source control
            Some(".gitattributes") | Some(".gitignore") | Some(".git") => continue,

            // Temporary Cargo files
            Some(".cargo-ok") => continue,

            // Skip patch-style orig/rej files. Published crates on crates.io
            // have `Cargo.toml.orig` which we don't want to use here and
            // otherwise these are rarely used as part of the build process.
            Some(filename) => {
                if filename.ends_with(".orig") || filename.ends_with(".rej") {
                    continue;
                }
            }
            _ => {}
        };

        // Join pathname components individually to make sure that the joined
        // path uses the correct directory separators everywhere, since
        // `relative` may use Unix-style and `dst` may require Windows-style
        // backslashes.
        let dst = relative
            .iter()
            .fold(dst.to_owned(), |acc, component| acc.join(&component));

        paths::create_dir_all(dst.parent().unwrap())?;

        let cksum = copy_and_checksum(p, &dst, tmp_buf)?;
        cksums.insert(relative.to_str().unwrap().replace("\\", "/"), cksum);
    }
    Ok(())
}

fn copy_and_checksum(src_path: &Path, dst_path: &Path, buf: &mut [u8]) -> CargoResult<String> {
    let mut src = File::open(src_path).with_context(|| format!("failed to open {:?}", src_path))?;
    let mut dst_opts = OpenOptions::new();
    dst_opts.write(true).create(true).truncate(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::{MetadataExt, OpenOptionsExt};
        let src_metadata = src
            .metadata()
            .with_context(|| format!("failed to stat {:?}", src_path))?;
        dst_opts.mode(src_metadata.mode());
    }
    let mut dst = dst_opts
        .open(dst_path)
        .with_context(|| format!("failed to create {:?}", dst_path))?;
    // Not going to bother setting mode on pre-existing files, since there
    // shouldn't be any under normal conditions.
    let mut cksum = Sha256::new();
    loop {
        let n = src
            .read(buf)
            .with_context(|| format!("failed to read from {:?}", src_path))?;
        if n == 0 {
            break Ok(cksum.finish_hex());
        }
        let data = &buf[..n];
        cksum.update(data);
        dst.write_all(data)
            .with_context(|| format!("failed to write to {:?}", dst_path))?;
    }
}
<commit_msg>Explain why normalizing slashes on vendor config output<commit_after>use crate::core::shell::Verbosity;
use crate::core::{GitReference, Workspace};
use crate::ops;
use crate::sources::path::PathSource;
use crate::sources::CRATES_IO_REGISTRY;
use crate::util::{CargoResult, Config};
use anyhow::{bail, Context as _};
use cargo_util::{paths, Sha256};
use serde::Serialize;
use std::collections::HashSet;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use toml_edit::easy as toml;

pub struct VendorOptions<'a> {
    pub no_delete: bool,
    pub versioned_dirs: bool,
    pub destination: &'a Path,
    pub extra: Vec<PathBuf>,
}

pub fn vendor(ws: &Workspace<'_>, opts: &VendorOptions<'_>) -> CargoResult<()> {
    let config = ws.config();
    let mut extra_workspaces = Vec::new();
    for extra in opts.extra.iter() {
        let extra = config.cwd().join(extra);
        let ws = Workspace::new(&extra, config)?;
        extra_workspaces.push(ws);
    }
    let workspaces = extra_workspaces.iter().chain(Some(ws)).collect::<Vec<_>>();
    let vendor_config = sync(config, &workspaces, opts).with_context(|| "failed to sync")?;

    if config.shell().verbosity() != Verbosity::Quiet {
        if vendor_config.source.is_empty() {
            crate::drop_eprintln!(config, "There is no dependency to vendor in this project.");
        } else {
            crate::drop_eprint!(
                config,
                "To use vendored sources, add this to your .cargo/config.toml for this project:\n\n"
            );
            crate::drop_print!(
                config,
                "{}",
                &toml::to_string_pretty(&vendor_config).unwrap()
            );
        }
    }

    Ok(())
}

#[derive(Serialize)]
struct VendorConfig {
    source: BTreeMap<String, VendorSource>,
}

#[derive(Serialize)]
#[serde(rename_all = "lowercase", untagged)]
enum VendorSource {
    Directory {
        directory: String,
    },
    Registry {
        registry: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
    Git {
        git: String,
        branch: Option<String>,
        tag: Option<String>,
        rev: Option<String>,
        #[serde(rename = "replace-with")]
        replace_with: String,
    },
}

fn sync(
    config: &Config,
    workspaces: &[&Workspace<'_>],
    opts: &VendorOptions<'_>,
) -> CargoResult<VendorConfig> {
    let canonical_destination = opts.destination.canonicalize();
    let canonical_destination = canonical_destination.as_deref().unwrap_or(opts.destination);
    let dest_dir_already_exists = canonical_destination.exists();

    paths::create_dir_all(&canonical_destination)?;
    let mut to_remove = HashSet::new();
    if !opts.no_delete {
        for entry in canonical_destination.read_dir()? {
            let entry = entry?;
            if !entry
                .file_name()
                .to_str()
                .map_or(false, |s| s.starts_with('.'))
            {
                to_remove.insert(entry.path());
            }
        }
    }

    // First up attempt to work around rust-lang/cargo#5956. Apparently build
    // artifacts sprout up in Cargo's global cache for whatever reason, although
    // it's unsure what tool is causing these issues at this time. For now we
    // apply a heavy-hammer approach which is to delete Cargo's unpacked version
    // of each crate to start off with. After we do this we'll re-resolve and
    // redownload again, which should trigger Cargo to re-extract all the
    // crates.
    //
    // Note that errors are largely ignored here as this is a best-effort
    // attempt. If anything fails here we basically just move on to the next
    // crate to work with.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).with_context(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .with_context(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // Don't delete actual source code!
            if pkg.source_id().is_path() {
                if let Ok(path) = pkg.source_id().url().to_file_path() {
                    if let Ok(path) = path.canonicalize() {
                        to_remove.remove(&path);
                    }
                }
                continue;
            }
            if pkg.source_id().is_git() {
                continue;
            }
            if let Ok(pkg) = packages.get_one(pkg) {
                drop(fs::remove_dir_all(pkg.manifest_path().parent().unwrap()));
            }
        }
    }

    let mut checksums = HashMap::new();
    let mut ids = BTreeMap::new();

    // Next up let's actually download all crates and start storing internal
    // tables about them.
    for ws in workspaces {
        let (packages, resolve) =
            ops::resolve_ws(ws).with_context(|| "failed to load pkg lockfile")?;

        packages
            .get_many(resolve.iter())
            .with_context(|| "failed to download packages")?;

        for pkg in resolve.iter() {
            // No need to vendor path crates since they're already in the
            // repository
            if pkg.source_id().is_path() {
                continue;
            }
            ids.insert(
                pkg,
                packages
                    .get_one(pkg)
                    .with_context(|| "failed to fetch package")?
                    .clone(),
            );

            checksums.insert(pkg, resolve.checksums().get(&pkg).cloned());
        }
    }

    let mut versions = HashMap::new();
    for id in ids.keys() {
        let map = versions.entry(id.name()).or_insert_with(BTreeMap::default);
        if let Some(prev) = map.get(&id.version()) {
            bail!(
                "found duplicate version of package `{} v{}` \
                 vendored from two sources:\n\
                 \n\
                 \tsource 1: {}\n\
                 \tsource 2: {}",
                id.name(),
                id.version(),
                prev,
                id.source_id()
            );
        }
        map.insert(id.version(), id.source_id());
    }

    let mut sources = BTreeSet::new();
    let mut tmp_buf = [0; 64 * 1024];
    for (id, pkg) in ids.iter() {
        // Next up, copy it to the vendor directory
        let src = pkg
            .manifest_path()
            .parent()
            .expect("manifest_path should point to a file");
        let max_version = *versions[&id.name()].iter().rev().next().unwrap().0;
        let dir_has_version_suffix = opts.versioned_dirs || id.version() != max_version;
        let dst_name = if dir_has_version_suffix {
            // Eg vendor/futures-0.1.13
            format!("{}-{}", id.name(), id.version())
        } else {
            // Eg vendor/futures
            id.name().to_string()
        };

        sources.insert(id.source_id());
        let dst = canonical_destination.join(&dst_name);
        to_remove.remove(&dst);
        let cksum = dst.join(".cargo-checksum.json");
        if dir_has_version_suffix && cksum.exists() {
            // Always re-copy directory without version suffix in case the version changed
            continue;
        }

        config.shell().status(
            "Vendoring",
            &format!("{} ({}) to {}", id, src.to_string_lossy(), dst.display()),
        )?;

        let _ = fs::remove_dir_all(&dst);
        let pathsource = PathSource::new(src, id.source_id(), config);
        let paths = pathsource.list_files(pkg)?;
        let mut map = BTreeMap::new();
        cp_sources(src, &paths, &dst, &mut map, &mut tmp_buf)
            .with_context(|| format!("failed to copy over vendored sources for: {}", id))?;

        // Finally, emit the metadata about this package
        let json = serde_json::json!({
            "package": checksums.get(id),
            "files": map,
        });

        paths::write(&cksum, json.to_string())?;
    }

    for path in to_remove {
        if path.is_dir() {
            paths::remove_dir_all(&path)?;
        } else {
            paths::remove_file(&path)?;
        }
    }

    // add our vendored source
    let mut config = BTreeMap::new();

    let merged_source_name = "vendored-sources";

    // replace original sources with vendor
    for source_id in sources {
        let name = if source_id.is_default_registry() {
            CRATES_IO_REGISTRY.to_string()
        } else {
            source_id.url().to_string()
        };

        let source = if source_id.is_default_registry() {
            VendorSource::Registry {
                registry: None,
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_remote_registry() {
            let registry = source_id.url().to_string();
            VendorSource::Registry {
                registry: Some(registry),
                replace_with: merged_source_name.to_string(),
            }
        } else if source_id.is_git() {
            let mut branch = None;
            let mut tag = None;
            let mut rev = None;
            if let Some(reference) = source_id.git_reference() {
                match *reference {
                    GitReference::Branch(ref b) => branch = Some(b.clone()),
                    GitReference::Tag(ref t) => tag = Some(t.clone()),
                    GitReference::Rev(ref r) => rev = Some(r.clone()),
                    GitReference::DefaultBranch => {}
                }
            }
            VendorSource::Git {
                git: source_id.url().to_string(),
                branch,
                tag,
                rev,
                replace_with: merged_source_name.to_string(),
            }
        } else {
            panic!("Invalid source ID: {}", source_id)
        };
        config.insert(name, source);
    }

    if !config.is_empty() {
        config.insert(
            merged_source_name.to_string(),
            VendorSource::Directory {
                // Windows-flavour paths are valid here on Windows but Unix.
                // This backslash normalization is for making output paths more
                // cross-platform compatible.
                directory: opts.destination.to_string_lossy().replace("\\", "/"),
            },
        );
    } else if !dest_dir_already_exists {
        // Nothing to vendor. Remove the destination dir we've just created.
        paths::remove_dir(canonical_destination)?;
    }

    Ok(VendorConfig { source: config })
}

fn cp_sources(
    src: &Path,
    paths: &[PathBuf],
    dst: &Path,
    cksums: &mut BTreeMap<String, String>,
    tmp_buf: &mut [u8],
) -> CargoResult<()> {
    for p in paths {
        let relative = p.strip_prefix(&src).unwrap();

        match relative.to_str() {
            // Skip git config files as they're not relevant to builds most of
            // the time and if we respect them (e.g.  in git) then it'll
            // probably mess with the checksums when a vendor dir is checked
            // into someone else's source control
            Some(".gitattributes") | Some(".gitignore") | Some(".git") => continue,

            // Temporary Cargo files
            Some(".cargo-ok") => continue,

            // Skip patch-style orig/rej files. Published crates on crates.io
            // have `Cargo.toml.orig` which we don't want to use here and
            // otherwise these are rarely used as part of the build process.
            Some(filename) => {
                if filename.ends_with(".orig") || filename.ends_with(".rej") {
                    continue;
                }
            }
            _ => {}
        };

        // Join pathname components individually to make sure that the joined
        // path uses the correct directory separators everywhere, since
        // `relative` may use Unix-style and `dst` may require Windows-style
        // backslashes.
        let dst = relative
            .iter()
            .fold(dst.to_owned(), |acc, component| acc.join(&component));

        paths::create_dir_all(dst.parent().unwrap())?;

        let cksum = copy_and_checksum(p, &dst, tmp_buf)?;
        cksums.insert(relative.to_str().unwrap().replace("\\", "/"), cksum);
    }
    Ok(())
}

fn copy_and_checksum(src_path: &Path, dst_path: &Path, buf: &mut [u8]) -> CargoResult<String> {
    let mut src = File::open(src_path).with_context(|| format!("failed to open {:?}", src_path))?;
    let mut dst_opts = OpenOptions::new();
    dst_opts.write(true).create(true).truncate(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::{MetadataExt, OpenOptionsExt};
        let src_metadata = src
            .metadata()
            .with_context(|| format!("failed to stat {:?}", src_path))?;
        dst_opts.mode(src_metadata.mode());
    }
    let mut dst = dst_opts
        .open(dst_path)
        .with_context(|| format!("failed to create {:?}", dst_path))?;
    // Not going to bother setting mode on pre-existing files, since there
    // shouldn't be any under normal conditions.
    let mut cksum = Sha256::new();
    loop {
        let n = src
            .read(buf)
            .with_context(|| format!("failed to read from {:?}", src_path))?;
        if n == 0 {
            break Ok(cksum.finish_hex());
        }
        let data = &buf[..n];
        cksum.update(data);
        dst.write_all(data)
            .with_context(|| format!("failed to write to {:?}", dst_path))?;
    }
}
<|endoftext|>


<commit_before>#![feature(test)]
#[cfg(test)]

extern crate jmespath;
extern crate test;
extern crate rustc_serialize;

use jmespath::Variable;
use std::rc::Rc;
use rustc_serialize::json::Json;
use test::Bencher;

#[bench]
fn bench_parsing_foo_bar_baz(b: &mut Bencher) {
    b.iter(|| jmespath::Expression::new("foo.bar.baz"));
}

#[bench]
fn bench_lexing_foo_bar_baz(b: &mut Bencher) {
    b.iter(|| for _ in jmespath::tokenize("foo.bar.baz") {});
}

#[bench]
fn bench_parse_and_interpret_foo_bar_baz(b: &mut Bencher) {
    let json = Json::from_str("{\"foo\":{\"bar\":{\"baz\":true}}}").unwrap();
    let var = Rc::new(Variable::from_json(&json));
    b.iter(|| {
        let expr = jmespath::Expression::new("foo.bar.baz").unwrap();
        expr.search(var.clone()).unwrap();
    });
}

#[bench]
fn bench_full_or_branches(b: &mut Bencher) {
    let json = Json::from_str("{\"foo\":true}").unwrap();
    let variable = Rc::new(Variable::from_json(&json));
    b.iter(|| {
        let expr = jmespath::Expression::new("bar || baz || bam || foo").unwrap();
        expr.search(variable.clone()).unwrap();
    });
}

#[bench]
fn bench_expr_identifier(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_expr_subexpr(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz.abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_deep_projection_104(b: &mut Bencher) {
    let deep = "a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*].a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*].a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*].a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*]";
    b.iter(|| jmespath::parse(deep));
}

#[bench]
fn bench_parse_identifier(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_subexpr(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz.abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_nested50(b: &mut Bencher) {
    let expr = "j49.j48.j47.j46.j45.j44.j43.j42.j41.j40.j39.j38.j37.j36.j35.j34.j33.j32.j31.j30.j29.j28.j27.j26.j25.j24.j23.j22.j21.j20.j19.j18.j17.j16.j15.j14.j13.j12.j11.j10.j9.j8.j7.j6.j5.j4.j3.j2.j1.j0";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_nested50_pipe(b: &mut Bencher) {
    let expr = "j49|j48|j47|j46|j45|j44|j43|j42|j41|j40|j39|j38|j37|j36|j35|j34|j33|j32|j31|j30|j29|j28|j27|j26|j25|j24|j23|j22|j21|j20|j19|j18|j17|j16|j15|j14|j13|j12|j11|j10|j9|j8|j7|j6|j5|j4|j3|j2|j1|j0";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_nested50_index(b: &mut Bencher) {
    let expr = "[49][48][47][46][45][44][43][42][41][40][39][38][37][36][35][34][33][32][31][30][29][28][27][26][25][24][23][22][21][20][19][18][17][16][15][14][13][12][11][10][9][8][7][6][5][4][3][2][1][0]";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_raw_string_literal(b: &mut Bencher) {
    let expr = "'abcdefghijklmnopqrstuvwxyz.abcdefghijklmnopqrstuvwxyz'";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_function_five_args(b: &mut Bencher) {
    let expr = "function(foo, bar, baz, bam, qux)";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_mutli_list_five_exprs(b: &mut Bencher) {
    let expr = "foo.[foo, bar, baz, bam, qux]";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_filter_projection(b: &mut Bencher) {
    let expr = "foo[bar > baz].qux";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_and_interpret_seven_deep_subexpr(b: &mut Bencher) {
    let expr = "a.b.c.d.e.f.g";
    let data = Json::from_str("{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":true}}}}}}}")
        .unwrap();
    let expression = jmespath::Expression::new(expr).unwrap();
    let var = Rc::new(Variable::from_json(&data));
    b.iter(|| expression.search(var.clone()));
}

#[bench]
fn bench_evaluate_deep(b: &mut Bencher) {
    let expr = "a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z.a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z.a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z.a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z";
    let expression = jmespath::Expression::new(expr).unwrap();
    let json = Json::from_str("{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":1}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}").unwrap();
    let variable = Rc::new(jmespath::Variable::from_json(&json));
    b.iter(|| expression.search(variable.clone()));
}
<commit_msg>Fixing benchmarks<commit_after>#![feature(test)]
#[cfg(test)]

extern crate jmespath;
extern crate test;
extern crate rustc_serialize;

use jmespath::Variable;
use jmespath::lexer::tokenize;
use std::rc::Rc;
use rustc_serialize::json::Json;
use test::Bencher;

#[bench]
fn bench_parsing_foo_bar_baz(b: &mut Bencher) {
    b.iter(|| jmespath::Expression::new("foo.bar.baz"));
}

#[bench]
fn bench_lexing_foo_bar_baz(b: &mut Bencher) {
    b.iter(|| for _ in tokenize("foo.bar.baz") {});
}

#[bench]
fn bench_parse_and_interpret_foo_bar_baz(b: &mut Bencher) {
    let json = Json::from_str("{\"foo\":{\"bar\":{\"baz\":true}}}").unwrap();
    let var = Rc::new(Variable::from_json(&json));
    b.iter(|| {
        let expr = jmespath::Expression::new("foo.bar.baz").unwrap();
        expr.search(var.clone()).unwrap();
    });
}

#[bench]
fn bench_full_or_branches(b: &mut Bencher) {
    let json = Json::from_str("{\"foo\":true}").unwrap();
    let variable = Rc::new(Variable::from_json(&json));
    b.iter(|| {
        let expr = jmespath::Expression::new("bar || baz || bam || foo").unwrap();
        expr.search(variable.clone()).unwrap();
    });
}

#[bench]
fn bench_expr_identifier(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_expr_subexpr(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz.abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_deep_projection_104(b: &mut Bencher) {
    let deep = "a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*].a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*].a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*].a[*].b[*].c[*].d[*].e[*].f[*].g[*].h[*].i[*].j[*].k[*].l[*].m[*].n[*].o[*].p[*].q[*].r[*].s[*].t[*].u[*].v[*].w[*].x[*].y[*].z[*]";
    b.iter(|| jmespath::parse(deep));
}

#[bench]
fn bench_parse_identifier(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_subexpr(b: &mut Bencher) {
    let expr = "abcdefghijklmnopqrstuvwxyz.abcdefghijklmnopqrstuvwxyz";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_nested50(b: &mut Bencher) {
    let expr = "j49.j48.j47.j46.j45.j44.j43.j42.j41.j40.j39.j38.j37.j36.j35.j34.j33.j32.j31.j30.j29.j28.j27.j26.j25.j24.j23.j22.j21.j20.j19.j18.j17.j16.j15.j14.j13.j12.j11.j10.j9.j8.j7.j6.j5.j4.j3.j2.j1.j0";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_nested50_pipe(b: &mut Bencher) {
    let expr = "j49|j48|j47|j46|j45|j44|j43|j42|j41|j40|j39|j38|j37|j36|j35|j34|j33|j32|j31|j30|j29|j28|j27|j26|j25|j24|j23|j22|j21|j20|j19|j18|j17|j16|j15|j14|j13|j12|j11|j10|j9|j8|j7|j6|j5|j4|j3|j2|j1|j0";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_nested50_index(b: &mut Bencher) {
    let expr = "[49][48][47][46][45][44][43][42][41][40][39][38][37][36][35][34][33][32][31][30][29][28][27][26][25][24][23][22][21][20][19][18][17][16][15][14][13][12][11][10][9][8][7][6][5][4][3][2][1][0]";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_raw_string_literal(b: &mut Bencher) {
    let expr = "'abcdefghijklmnopqrstuvwxyz.abcdefghijklmnopqrstuvwxyz'";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_function_five_args(b: &mut Bencher) {
    let expr = "function(foo, bar, baz, bam, qux)";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_mutli_list_five_exprs(b: &mut Bencher) {
    let expr = "foo.[foo, bar, baz, bam, qux]";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_filter_projection(b: &mut Bencher) {
    let expr = "foo[bar > baz].qux";
    b.iter(|| jmespath::parse(expr));
}

#[bench]
fn bench_parse_and_interpret_seven_deep_subexpr(b: &mut Bencher) {
    let expr = "a.b.c.d.e.f.g";
    let data = Json::from_str("{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":true}}}}}}}")
        .unwrap();
    let expression = jmespath::Expression::new(expr).unwrap();
    let var = Rc::new(Variable::from_json(&data));
    b.iter(|| expression.search(var.clone()));
}

#[bench]
fn bench_evaluate_deep(b: &mut Bencher) {
    let expr = "a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z.a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z.a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z.a.b.c.d.e.f.g.h.i.j.k.l.m.n.o.p.q.r.s.t.u.v.w.x.y.z";
    let expression = jmespath::Expression::new(expr).unwrap();
    let json = Json::from_str("{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":{\"a\":{\"b\":{\"c\":{\"d\":{\"e\":{\"f\":{\"g\":{\"h\":{\"i\":{\"j\":{\"k\":{\"l\":{\"m\":{\"n\":{\"o\":{\"p\":{\"q\":{\"r\":{\"s\":{\"t\":{\"u\":{\"v\":{\"w\":{\"x\":{\"y\":{\"z\":1}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}").unwrap();
    let variable = Rc::new(jmespath::Variable::from_json(&json));
    b.iter(|| expression.search(variable.clone()));
}
<|endoftext|>


<commit_before>/// A genetic algorithm in Rust
/// Copyright (C) 2015  Andrew Schwartzmeyer

use std::cmp::{Eq, PartialEq, Ordering, PartialOrd};
use std::mem;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};

/// Schwefel's Evolutionary Computation benchmark [problem][]
/// Evaluated on the hypercube with x_i in [-500, 500] and dim(x) = 50
/// Global minimum: f(x*) = 0, at x* = (420.9687, ...)
/// [problem]: http://www.sfu.ca/~ssurjano/schwef.html
fn schwefel(solution: &[f64]) -> f64 {
    418.9829_f64 * solution.len() as f64
        - solution.iter().fold(0_f64, |sum, x| {
            sum + x * x.abs().sqrt().sin()
        })
}

/// An Orderable, Cloneable solution with a cached fitness
pub struct Individual {
    pub solution: Vec<f64>,
    pub fitness: f64,
}

impl Individual {
    /// Constructs a new Individual with 50 random values in the range
    pub fn new<R: Rng>(range: &Range<f64>, rng: &mut R) -> Self {
        let solution: Vec<_> = (0..50).map(|_| range.ind_sample(rng)).collect();
        let fitness = schwefel(&solution);
        Individual { solution: solution, fitness: fitness }
    }

    /// Mutates 1 in 40 genes to a new value in the range
    pub fn mutate<R: Rng>(&mut self, range: &Range<f64>, rng: &mut R) {
        for x in &mut self.solution {
            if rng.gen_weighted_bool(40) {
                *x = range.ind_sample(rng)
            }
        }
        self.fitness = schwefel(&self.solution);
    }

    /// Recombines two Individuals with a 1 in 2 chance via two-point crossover
    pub fn combine<R: Rng>(x: &mut Individual, y: &mut Individual, rng: &mut R) {
        if rng.gen_weighted_bool(2) {
            let len = x.solution.len();
            let (begin, n) = (rng.gen_range(0, len), rng.gen_range(0, len));
            for i in begin..(begin + n) % len {
                mem::swap(&mut x.solution[i], &mut y.solution[i]);
            }
            x.fitness = schwefel(&x.solution);
            y.fitness = schwefel(&y.solution);
        }
    }
}

impl Clone for Individual {
    fn clone(&self) -> Self {
        Individual { solution: self.solution.clone(), fitness: self.fitness }
    }
}

impl Eq for Individual {}

impl Ord for Individual {
    /// This dangerously delegates to `partial_cmp`; NaN will panic
    fn cmp(&self, other: &Self) -> Ordering {
        if let Some(result) = self.fitness.partial_cmp(&other.fitness) {
            return result;
        }
        unimplemented!();
    }
}

impl PartialEq<Individual> for Individual {
    /// This doesn't use `fitness.eq()` because it needs to be
    /// consistent with `Eq`
    fn eq(&self, other: &Individual) -> bool {
        if let Some(result) = self.fitness.partial_cmp(&other.fitness) {
            return result == Ordering::Equal;
        }
        unimplemented!();
    }
}

impl PartialOrd for Individual {
    /// This delegates to `fitness.partial_cmp()`
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.fitness.partial_cmp(&other.fitness)
    }
}
<commit_msg>Remove rand bottleneck from mutate()<commit_after>/// A genetic algorithm in Rust
/// Copyright (C) 2015  Andrew Schwartzmeyer

use std::cmp::{Eq, PartialEq, Ordering, PartialOrd};
use std::mem;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};

/// Schwefel's Evolutionary Computation benchmark [problem][]
/// Evaluated on the hypercube with x_i in [-500, 500] and dim(x) = 50
/// Global minimum: f(x*) = 0, at x* = (420.9687, ...)
/// [problem]: http://www.sfu.ca/~ssurjano/schwef.html
fn schwefel(solution: &[f64]) -> f64 {
    418.9829_f64 * solution.len() as f64
        - solution.iter().fold(0_f64, |sum, x| {
            sum + x * x.abs().sqrt().sin()
        })
}

/// An Orderable, Cloneable solution with a cached fitness
pub struct Individual {
    pub solution: Vec<f64>,
    pub fitness: f64,
}

impl Individual {
    /// Constructs a new Individual with 50 random values in the range
    pub fn new<R: Rng>(range: &Range<f64>, rng: &mut R) -> Self {
        let solution: Vec<_> = (0..50).map(|_| range.ind_sample(rng)).collect();
        let fitness = schwefel(&solution);
        Individual { solution: solution, fitness: fitness }
    }

    /// Mutates 0 or 1 random genes to a new value in the range
    pub fn mutate<R: Rng>(&mut self, range: &Range<f64>, rng: &mut R) {
        for _ in 0..rng.gen_range(0, 2) {
            let i = rng.gen_range(0, self.solution.len());
            self.solution[i] = range.ind_sample(rng);
        }
        self.fitness = schwefel(&self.solution);
    }

    /// Recombines two Individuals with a 1 in 2 chance via two-point crossover
    pub fn combine<R: Rng>(x: &mut Individual, y: &mut Individual, rng: &mut R) {
        if rng.gen_weighted_bool(2) {
            let len = x.solution.len();
            let (begin, n) = (rng.gen_range(0, len), rng.gen_range(0, len));
            for i in begin..(begin + n) % len {
                mem::swap(&mut x.solution[i], &mut y.solution[i]);
            }
            x.fitness = schwefel(&x.solution);
            y.fitness = schwefel(&y.solution);
        }
    }
}

impl Clone for Individual {
    fn clone(&self) -> Self {
        Individual { solution: self.solution.clone(), fitness: self.fitness }
    }
}

impl Eq for Individual {}

impl Ord for Individual {
    /// This dangerously delegates to `partial_cmp`; NaN will panic
    fn cmp(&self, other: &Self) -> Ordering {
        if let Some(result) = self.fitness.partial_cmp(&other.fitness) {
            return result;
        }
        unimplemented!();
    }
}

impl PartialEq<Individual> for Individual {
    /// This doesn't use `fitness.eq()` because it needs to be
    /// consistent with `Eq`
    fn eq(&self, other: &Individual) -> bool {
        if let Some(result) = self.fitness.partial_cmp(&other.fitness) {
            return result == Ordering::Equal;
        }
        unimplemented!();
    }
}

impl PartialOrd for Individual {
    /// This delegates to `fitness.partial_cmp()`
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.fitness.partial_cmp(&other.fitness)
    }
}
<|endoftext|>


<commit_before>
type InterruptHandlerFn = unsafe extern "C" fn() -> !;

use dispatch_table::{DispatchTable, DispatchFn};
use spin::RwLock;
use vmx;
use core::mem;

#[derive(Default)]
#[repr(packed)]
struct IDTEntry {
    base_low: u16,
    selector: u16,
    always0: u8,
    flags: u8,
    base_high: u16,
    base_highest: u32,
    _reserved: u32,
}

fn new_host_idt_descriptor() -> vmx::CPUTableDescriptor {
    let base = (*DESCRIPTOR_TABLE.write()).0.as_ptr() as u64;
    vmx::CPUTableDescriptor{
        base: base,
        limit: (mem::size_of::<IDT>() - 1) as u16,
    }
}

#[derive(Default)]
pub struct IDT([IDTEntry; 20]);

lazy_static! {
    static ref INTERRUPT_TABLE: RwLock<DispatchTable<u64>> = RwLock::new(DispatchTable::new(256));
    static ref DESCRIPTOR_TABLE: RwLock<IDT> = RwLock::new(Default::default());
}

impl IDT {
    fn set_entry(&mut self, num: usize, base: u64, selector: u16, flags: u8) {
        self.0[num].base_low = (base & 0xffff) as u16;
        self.0[num].base_high = ((base >> 16) & 0xffff) as u16;
        self.0[num].base_highest = (base >> 32) as u32;
        self.0[num].always0 = 0;
        self.0[num].selector = selector;
        self.0[num].flags = flags;
    }
}

#[repr(C)]
pub struct InterruptCPUState {
	ds: u64,
	r15: u64,
	r14: u64,
	r13: u64,
	r12: u64,
	r11: u64,
	r10: u64,
	r9: u64,
	r8: u64,
	rdi: u64,
	rsi: u64,
	rbp: u64,
	rdx: u64,
	rcx: u64,
	rbx: u64,
	rax: u64,
	interrupt_number: u64,
	error_code: u64,
	rip: u64,
	cs: u64,
	rflags: u64,
	rsp: u64,
	ss: u64,
}

pub fn register_interrupt_handler(interrupt: u64, stub: InterruptHandlerFn, handler: DispatchFn<u64>) {
    INTERRUPT_TABLE.write().register(interrupt, handler);
    DESCRIPTOR_TABLE.write().set_entry(interrupt as usize, stub as u64, 0x08, 0x8e);
}

pub fn unregister_interrupt_handler(interrupt: u64) {
    INTERRUPT_TABLE.write().unregister(interrupt);
}

#[no_mangle]
pub extern "C" fn dispatch_interrupt(state: &mut InterruptCPUState) {
    INTERRUPT_TABLE.read().dispatch(&state.interrupt_number);
}

fn default_handler(_interrupt_number: &u64) -> bool {
    true
}

extern "C" {
    fn _isr0() -> !;
    fn _isr1() -> !;
    fn _isr2() -> !;
    fn _isr3() -> !;
    fn _isr4() -> !;
    fn _isr5() -> !;
    fn _isr6() -> !;
    fn _isr7() -> !;
    fn _isr8() -> !;
    fn _isr9() -> !;
    fn _isr10() -> !;
    fn _isr11() -> !;
    fn _isr12() -> !;
    fn _isr13() -> !;
    fn _isr14() -> !;
    fn _isr15() -> !;
    fn _isr16() -> !;
    fn _isr17() -> !;
    fn _isr18() -> !;
    fn _isr19() -> !;
}

const ISR: [InterruptHandlerFn; 20] = [_isr0, _isr1, _isr2, _isr3, _isr4, _isr5, _isr6, _isr7, _isr8, _isr9, _isr10, _isr11, _isr12, _isr13, _isr14, _isr15, _isr16, _isr17, _isr18, _isr19];

pub fn init_interrupt_handlers() {
    for i in 0..20 {
        register_interrupt_handler(i, ISR[i as usize], default_handler);
    }
}

#[cfg(feature = "runtime_tests")]
pub mod runtime_tests {

    use interrupts;
    use vmx;

}
<commit_msg>IDT runtime tests<commit_after>
type InterruptHandlerFn = unsafe extern "C" fn() -> !;

use dispatch_table::{DispatchTable, DispatchFn};
use spin::RwLock;
use vmx;
use core::mem;

#[derive(Default)]
#[repr(packed)]
struct IDTEntry {
    base_low: u16,
    selector: u16,
    always0: u8,
    flags: u8,
    base_high: u16,
    base_highest: u32,
    _reserved: u32,
}

fn new_host_idt_descriptor() -> vmx::CPUTableDescriptor {
    let base = (*DESCRIPTOR_TABLE.write()).0.as_ptr() as u64;
    vmx::CPUTableDescriptor{
        base: base,
        limit: (mem::size_of::<IDT>() - 1) as u16,
    }
}

#[derive(Default)]
pub struct IDT([IDTEntry; 20]);

lazy_static! {
    static ref INTERRUPT_TABLE: RwLock<DispatchTable<u64>> = RwLock::new(DispatchTable::new(256));
    static ref DESCRIPTOR_TABLE: RwLock<IDT> = RwLock::new(Default::default());
}

impl IDT {
    fn set_entry(&mut self, num: usize, base: u64, selector: u16, flags: u8) {
        self.0[num].base_low = (base & 0xffff) as u16;
        self.0[num].base_high = ((base >> 16) & 0xffff) as u16;
        self.0[num].base_highest = (base >> 32) as u32;
        self.0[num].always0 = 0;
        self.0[num].selector = selector;
        self.0[num].flags = flags;
    }
}

#[repr(C)]
pub struct InterruptCPUState {
	ds: u64,
	r15: u64,
	r14: u64,
	r13: u64,
	r12: u64,
	r11: u64,
	r10: u64,
	r9: u64,
	r8: u64,
	rdi: u64,
	rsi: u64,
	rbp: u64,
	rdx: u64,
	rcx: u64,
	rbx: u64,
	rax: u64,
	interrupt_number: u64,
	error_code: u64,
	rip: u64,
	cs: u64,
	rflags: u64,
	rsp: u64,
	ss: u64,
}

pub fn register_interrupt_handler(interrupt: u64, stub: InterruptHandlerFn, handler: DispatchFn<u64>) {
    INTERRUPT_TABLE.write().register(interrupt, handler);
    DESCRIPTOR_TABLE.write().set_entry(interrupt as usize, stub as u64, 0x08, 0x8e);
}

pub fn unregister_interrupt_handler(interrupt: u64) {
    INTERRUPT_TABLE.write().unregister(interrupt);
}

#[no_mangle]
pub extern "C" fn dispatch_interrupt(state: &mut InterruptCPUState) {
    INTERRUPT_TABLE.read().dispatch(&state.interrupt_number);
}

fn default_handler(_interrupt_number: &u64) -> bool {
    true
}

extern "C" {
    fn _isr0() -> !;
    fn _isr1() -> !;
    fn _isr2() -> !;
    fn _isr3() -> !;
    fn _isr4() -> !;
    fn _isr5() -> !;
    fn _isr6() -> !;
    fn _isr7() -> !;
    fn _isr8() -> !;
    fn _isr9() -> !;
    fn _isr10() -> !;
    fn _isr11() -> !;
    fn _isr12() -> !;
    fn _isr13() -> !;
    fn _isr14() -> !;
    fn _isr15() -> !;
    fn _isr16() -> !;
    fn _isr17() -> !;
    fn _isr18() -> !;
    fn _isr19() -> !;
}

const ISR: [InterruptHandlerFn; 20] = [_isr0, _isr1, _isr2, _isr3, _isr4, _isr5, _isr6, _isr7, _isr8, _isr9, _isr10, _isr11, _isr12, _isr13, _isr14, _isr15, _isr16, _isr17, _isr18, _isr19];

pub fn init_interrupt_handlers() {
    for i in 0..20 {
        register_interrupt_handler(i, ISR[i as usize], default_handler);
    }
}

#[cfg(feature = "runtime_tests")]
pub mod runtime_tests {

    use interrupts;
    use vmx;
    use cli;

    pub fn run() {
        test_load_and_restore_idt();
    }

    fn test_load_and_restore_idt() {
        cli::ClearLocalInterruptsGuard::new();
        let idt_desc = interrupts::new_host_idt_descriptor();
        let mut orig_idt_desc: vmx::CPUTableDescriptor = Default::default();
        vmx::sidt(&mut orig_idt_desc);
        vmx::lidt(&idt_desc);
        vmx::lidt(&orig_idt_desc);
    }

    /*
    fn division_by_zero_handler() {

    }

    fn test_divide_by_zero_interrupt() {
        let mut orig_idt_desc: vmx::CPUTableDescriptor = Default::default();

        {
            cli::ClearLocalInterruptsGuard::new();
            let idt_desc = interrupts::new_host_idt_descriptor();
            vmx::lidt(&idt_desc);
        }

        let a = 10;
        let b = 0;
        let c = a / b; // Division by zero interrupt

    }
    */
}
<|endoftext|>


<commit_before>use libc::{c_char, c_void};

#[link(name = "dl")]
extern "C" {
    fn dlsym(handle: *const c_void, symbol: *const c_char) -> *const c_void;
}

const RTLD_NEXT: *const c_void = -1isize as *const c_void;

pub unsafe fn dlsym_next(symbol: &'static str) -> *const u8 {
    let ptr = dlsym(RTLD_NEXT, symbol.as_ptr() as *const c_char);
    if ptr.is_null() {
        panic!("redhook: Unable to find underlying function for {}", symbol);
    }
    ptr as *const u8
}

/* Rust doesn't directly expose __attribute__((constructor)), but this
 * is how GNU implements it. */
#[link_section = ".init_array"]
pub static INITIALIZE_CTOR: extern "C" fn() = ::initialize;

#[macro_export]
macro_rules! hook {
    (unsafe fn $real_fn:ident ( $($v:ident : $t:ty),* ) -> $r:ty => $hook_fn:ident $body:block) => {
        #[allow(non_camel_case_types)]
        pub struct $real_fn {__private_field: ()}
        #[allow(non_upper_case_globals)]
        static $real_fn: $real_fn = $real_fn {__private_field: ()};

        impl $real_fn {
            fn get(&self) -> unsafe extern fn ( $($v : $t),* ) -> $r {
                use ::std::sync::{Once, ONCE_INIT};

                static mut REAL: *const u8 = 0 as *const u8;
                static mut ONCE: Once = ONCE_INIT;

                unsafe {
                    ONCE.call_once(|| {
                        REAL = $crate::ld_preload::dlsym_next(concat!(stringify!($real_fn), "\0"));
                    });
                    ::std::mem::transmute(REAL)
                }
            }

            #[no_mangle]
            pub unsafe extern fn $real_fn ( $($v : $t),* ) -> $r {
                if $crate::initialized() {
                    ::std::panic::catch_unwind(|| $hook_fn ( $($v),* )).ok()
                } else {
                    None
                }.unwrap_or_else(|| $real_fn.get() ( $($v),* ))
            }
        }

        pub unsafe fn $hook_fn ( $($v : $t),* ) -> $r {
            $body
        }
    };

    (unsafe fn $real_fn:ident ( $($v:ident : $t:ty),* ) => $hook_fn:ident $body:block) => {
        hook! { unsafe fn $real_fn ( $($v : $t),* ) -> () => $hook_fn $body }
    };
}

#[macro_export]
macro_rules! real {
    ($real_fn:ident) => {
        $real_fn.get()
    };
}
<commit_msg>Fixes #11 -- correctly scope hook!<commit_after>use libc::{c_char, c_void};

#[link(name = "dl")]
extern "C" {
    fn dlsym(handle: *const c_void, symbol: *const c_char) -> *const c_void;
}

const RTLD_NEXT: *const c_void = -1isize as *const c_void;

pub unsafe fn dlsym_next(symbol: &'static str) -> *const u8 {
    let ptr = dlsym(RTLD_NEXT, symbol.as_ptr() as *const c_char);
    if ptr.is_null() {
        panic!("redhook: Unable to find underlying function for {}", symbol);
    }
    ptr as *const u8
}

/* Rust doesn't directly expose __attribute__((constructor)), but this
 * is how GNU implements it. */
#[link_section = ".init_array"]
pub static INITIALIZE_CTOR: extern "C" fn() = ::initialize;

#[macro_export]
macro_rules! hook {
    (unsafe fn $real_fn:ident ( $($v:ident : $t:ty),* ) -> $r:ty => $hook_fn:ident $body:block) => {
        #[allow(non_camel_case_types)]
        pub struct $real_fn {__private_field: ()}
        #[allow(non_upper_case_globals)]
        static $real_fn: $real_fn = $real_fn {__private_field: ()};

        impl $real_fn {
            fn get(&self) -> unsafe extern fn ( $($v : $t),* ) -> $r {
                use ::std::sync::{Once, ONCE_INIT};

                static mut REAL: *const u8 = 0 as *const u8;
                static mut ONCE: Once = ONCE_INIT;

                unsafe {
                    ONCE.call_once(|| {
                        REAL = $crate::ld_preload::dlsym_next(concat!(stringify!($real_fn), "\0"));
                    });
                    ::std::mem::transmute(REAL)
                }
            }

            #[no_mangle]
            pub unsafe extern fn $real_fn ( $($v : $t),* ) -> $r {
                if $crate::initialized() {
                    ::std::panic::catch_unwind(|| $hook_fn ( $($v),* )).ok()
                } else {
                    None
                }.unwrap_or_else(|| $real_fn.get() ( $($v),* ))
            }
        }

        pub unsafe fn $hook_fn ( $($v : $t),* ) -> $r {
            $body
        }
    };

    (unsafe fn $real_fn:ident ( $($v:ident : $t:ty),* ) => $hook_fn:ident $body:block) => {
        $crate::hook! { unsafe fn $real_fn ( $($v : $t),* ) -> () => $hook_fn $body }
    };
}

#[macro_export]
macro_rules! real {
    ($real_fn:ident) => {
        $real_fn.get()
    };
}
<|endoftext|>


<commit_before>// Copyright 2014 Michael Yang. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

//! Wrappers for matrix functions.

use num::complex::{Complex, Complex32, Complex64};
use attribute::{Diagonal, Side, Symmetry, Transpose};
use pointer::CPtr;
use scalar::Scalar;
use matrix::ll::*;
use matrix::Matrix;
use vector::Vector;

pub trait Gemm {
    fn gemm(alpha: &Self, at: Transpose, a: &Matrix<Self>, bt: Transpose, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

macro_rules! gemm_impl(($($t: ident), +) => (
    $(
        impl Gemm for $t {
            fn gemm(alpha: &$t, at: Transpose, a: &Matrix<$t>, bt: Transpose, b: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, gemm)(a.order(),
                        at, bt,
                        a.rows(), b.cols(), a.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

gemm_impl!(f32, f64, Complex32, Complex64);

#[cfg(test)]
mod gemm_tests {
    use std::iter::repeat;
    use attribute::Transpose;
    use matrix::ops::Gemm;

    #[test]
    fn real() {
        let a = (2, 2, vec![1.0, 2.0, 3.0, 4.0]);
        let b = (2, 2, vec![-1.0, 3.0, 1.0, 1.0]);
        let t = Transpose::NoTrans;

        let mut c = (2, 2, repeat(0.0).take(4).collect());
        Gemm::gemm(&1f32, t, &a, t, &b, &0f32, &mut c);

        assert_eq!(c.2, vec![1.0, 5.0, 1.0, 13.0]);
    }
}

pub trait Symm {
    fn symm(side: Side, symmetry: Symmetry, alpha: &Self, a: &Matrix<Self>, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

pub trait Hemm {
    fn hemm(side: Side, symmetry: Symmetry, alpha: &Self, a: &Matrix<Self>, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

macro_rules! symm_impl(($trait_name: ident, $fn_name: ident, $($t: ident), +) => (
    $(
        impl $trait_name for $t {
            fn $fn_name(side: Side, symmetry: Symmetry, alpha: &$t, a: &Matrix<$t>, b: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, $fn_name)(a.order(),
                        side, symmetry,
                        a.rows(), b.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

symm_impl!(Symm, symm, f32, f64, Complex32, Complex64);
symm_impl!(Hemm, hemm, Complex32, Complex64);

pub trait Trmm {
    fn trmm(side: Side, symmetry: Symmetry, trans: Transpose, diag: Diagonal, alpha: &Self, a: &Matrix<Self>, b: &mut Matrix<Self>);
}

pub trait Trsm {
    fn trsm(side: Side, symmetry: Symmetry, trans: Transpose, diag: Diagonal, alpha: &Self, a: &Matrix<Self>, b: &mut Matrix<Self>);
}

macro_rules! trmm_impl(($trait_name: ident, $fn_name: ident, $($t: ident), +) => (
    $(
        impl $trait_name for $t {
            fn $fn_name(side: Side, symmetry: Symmetry, trans: Transpose, diag: Diagonal, alpha: &$t, a: &Matrix<$t>, b: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, $fn_name)(a.order(),
                        side, symmetry, trans, diag,
                        b.rows(), b.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_mut_ptr().as_c_ptr(), b.lead_dim());
                }
            }
        }
    )+
));

trmm_impl!(Trmm, trmm, f32, f64, Complex32, Complex64);
trmm_impl!(Trsm, trsm, Complex32, Complex64);

pub trait Herk {
    fn herk(symmetry: Symmetry, trans: Transpose, alpha: &Self, a: &Matrix<Complex<Self>>, beta: &Self, c: &mut Matrix<Complex<Self>>);
}

pub trait Her2k {
    fn her2k(symmetry: Symmetry, trans: Transpose, alpha: Complex<Self>, a: &Matrix<Complex<Self>>, b: &Matrix<Complex<Self>>, beta: &Self, c: &mut Matrix<Complex<Self>>);
}

macro_rules! herk_impl(($($t: ident), +) => (
    $(
        impl Herk for $t {
            fn herk(symmetry: Symmetry, trans: Transpose, alpha: &$t, a: &Matrix<Complex<$t>>, beta: &$t, c: &mut Matrix<Complex<$t>>) {
                unsafe {
                    prefix!(Complex<$t>, herk)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        *alpha,
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        *beta,
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }

        impl Her2k for $t {
            fn her2k(symmetry: Symmetry, trans: Transpose, alpha: Complex<$t>, a: &Matrix<Complex<$t>>, b: &Matrix<Complex<$t>>, beta: &$t, c: &mut Matrix<Complex<$t>>) {
                unsafe {
                    prefix!(Complex<$t>, her2k)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        *beta,
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

herk_impl!(f32, f64);

pub trait Syrk {
    fn syrk(symmetry: Symmetry, trans: Transpose, alpha: &Self, a: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

pub trait Syr2k {
    fn syr2k(symmetry: Symmetry, trans: Transpose, alpha: &Self, a: &Matrix<Self>, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

macro_rules! syrk_impl(($($t: ident), +) => (
    $(
        impl Syrk for $t {
            fn syrk(symmetry: Symmetry, trans: Transpose, alpha: &$t, a: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, syrk)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }

        impl Syr2k for $t {
            fn syr2k(symmetry: Symmetry, trans: Transpose, alpha: &$t, a: &Matrix<$t>, b: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, syr2k)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

syrk_impl!(f32, f64, Complex32, Complex64);
<commit_msg>Fix gemm when one or both matrices is transposed.<commit_after>// Copyright 2014 Michael Yang. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

//! Wrappers for matrix functions.

use num::complex::{Complex, Complex32, Complex64};
use attribute::{Diagonal, Side, Symmetry, Transpose};
use pointer::CPtr;
use scalar::Scalar;
use matrix::ll::*;
use matrix::Matrix;
use vector::Vector;

pub trait Gemm {
    fn gemm(alpha: &Self, at: Transpose, a: &Matrix<Self>, bt: Transpose, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

macro_rules! gemm_impl(($($t: ident), +) => (
    $(
        impl Gemm for $t {
            fn gemm(alpha: &$t, at: Transpose, a: &Matrix<$t>, bt: Transpose, b: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    let (m, k)  = match at {
                        Transpose::NoTrans => (a.rows(), a.cols()),
                        _ => (a.cols(), a.rows()),
                    };

                    let n = match bt {
                        Transpose::NoTrans => b.cols(),
                        _ => b.rows(),
                    };

                    prefix!($t, gemm)(a.order(),
                        at, bt,
                        m, n, k,
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

gemm_impl!(f32, f64, Complex32, Complex64);

#[cfg(test)]
mod gemm_tests {
    use std::iter::repeat;
    use attribute::Transpose;
    use matrix::ops::Gemm;

    #[test]
    fn real() {
        let a = (2, 2, vec![1.0, 2.0, 3.0, 4.0]);
        let b = (2, 2, vec![-1.0, 3.0, 1.0, 1.0]);
        let t = Transpose::NoTrans;

        let mut c = (2, 2, repeat(0.0).take(4).collect());
        Gemm::gemm(&1f32, t, &a, t, &b, &0f32, &mut c);

        assert_eq!(c.2, vec![1.0, 5.0, 1.0, 13.0]);
    }

    #[test]
    fn transpose() {
        let a = (3, 2, vec![
                 1.0, 2.0,
                 3.0, 4.0,
                 5.0, 6.0, ]);
        let b = (2, 3, vec![
                 -1.0, 3.0, 1.0,
                 1.0, 1.0, 1.0]);
        let t = Transpose::Trans;

        let mut c = (2, 2, repeat(0.0).take(4).collect());
        Gemm::gemm(&1f32, t, &a, t, &b, &0f32, &mut c);

        assert_eq!(c.2, vec![13.0, 9.0, 16.0, 12.0]);
    }
}

pub trait Symm {
    fn symm(side: Side, symmetry: Symmetry, alpha: &Self, a: &Matrix<Self>, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

pub trait Hemm {
    fn hemm(side: Side, symmetry: Symmetry, alpha: &Self, a: &Matrix<Self>, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

macro_rules! symm_impl(($trait_name: ident, $fn_name: ident, $($t: ident), +) => (
    $(
        impl $trait_name for $t {
            fn $fn_name(side: Side, symmetry: Symmetry, alpha: &$t, a: &Matrix<$t>, b: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, $fn_name)(a.order(),
                        side, symmetry,
                        a.rows(), b.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

symm_impl!(Symm, symm, f32, f64, Complex32, Complex64);
symm_impl!(Hemm, hemm, Complex32, Complex64);

pub trait Trmm {
    fn trmm(side: Side, symmetry: Symmetry, trans: Transpose, diag: Diagonal, alpha: &Self, a: &Matrix<Self>, b: &mut Matrix<Self>);
}

pub trait Trsm {
    fn trsm(side: Side, symmetry: Symmetry, trans: Transpose, diag: Diagonal, alpha: &Self, a: &Matrix<Self>, b: &mut Matrix<Self>);
}

macro_rules! trmm_impl(($trait_name: ident, $fn_name: ident, $($t: ident), +) => (
    $(
        impl $trait_name for $t {
            fn $fn_name(side: Side, symmetry: Symmetry, trans: Transpose, diag: Diagonal, alpha: &$t, a: &Matrix<$t>, b: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, $fn_name)(a.order(),
                        side, symmetry, trans, diag,
                        b.rows(), b.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_mut_ptr().as_c_ptr(), b.lead_dim());
                }
            }
        }
    )+
));

trmm_impl!(Trmm, trmm, f32, f64, Complex32, Complex64);
trmm_impl!(Trsm, trsm, Complex32, Complex64);

pub trait Herk {
    fn herk(symmetry: Symmetry, trans: Transpose, alpha: &Self, a: &Matrix<Complex<Self>>, beta: &Self, c: &mut Matrix<Complex<Self>>);
}

pub trait Her2k {
    fn her2k(symmetry: Symmetry, trans: Transpose, alpha: Complex<Self>, a: &Matrix<Complex<Self>>, b: &Matrix<Complex<Self>>, beta: &Self, c: &mut Matrix<Complex<Self>>);
}

macro_rules! herk_impl(($($t: ident), +) => (
    $(
        impl Herk for $t {
            fn herk(symmetry: Symmetry, trans: Transpose, alpha: &$t, a: &Matrix<Complex<$t>>, beta: &$t, c: &mut Matrix<Complex<$t>>) {
                unsafe {
                    prefix!(Complex<$t>, herk)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        *alpha,
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        *beta,
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }

        impl Her2k for $t {
            fn her2k(symmetry: Symmetry, trans: Transpose, alpha: Complex<$t>, a: &Matrix<Complex<$t>>, b: &Matrix<Complex<$t>>, beta: &$t, c: &mut Matrix<Complex<$t>>) {
                unsafe {
                    prefix!(Complex<$t>, her2k)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        *beta,
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

herk_impl!(f32, f64);

pub trait Syrk {
    fn syrk(symmetry: Symmetry, trans: Transpose, alpha: &Self, a: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

pub trait Syr2k {
    fn syr2k(symmetry: Symmetry, trans: Transpose, alpha: &Self, a: &Matrix<Self>, b: &Matrix<Self>, beta: &Self, c: &mut Matrix<Self>);
}

macro_rules! syrk_impl(($($t: ident), +) => (
    $(
        impl Syrk for $t {
            fn syrk(symmetry: Symmetry, trans: Transpose, alpha: &$t, a: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, syrk)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }

        impl Syr2k for $t {
            fn syr2k(symmetry: Symmetry, trans: Transpose, alpha: &$t, a: &Matrix<$t>, b: &Matrix<$t>, beta: &$t, c: &mut Matrix<$t>) {
                unsafe {
                    prefix!($t, syr2k)(a.order(),
                        symmetry, trans,
                        a.rows(), a.cols(),
                        alpha.as_const(),
                        a.as_ptr().as_c_ptr(), a.lead_dim(),
                        b.as_ptr().as_c_ptr(), b.lead_dim(),
                        beta.as_const(),
                        c.as_mut_ptr().as_c_ptr(), c.lead_dim());
                }
            }
        }
    )+
));

syrk_impl!(f32, f64, Complex32, Complex64);
<|endoftext|>


<commit_before>use dns_parser::Packet as DnsPacket;
use dns_parser::Opcode;
use serde_json::{Value, Map, Number};

pub struct DnsHeader {
	id: Number,
	opcode: String,
	flags: Value,
	total_questions: Number,
	total_answer_rrs: Number,
	total_authority_rrs: Number,
	total_additional_rrs: Number,
}

impl DnsHeader {
	pub fn new(p: &DnsPacket) -> DnsHeader {
		let opcode_number = p.header.opcode;
		let opcode_string = match opcode_number {
			Opcode::StandardQuery => "QUERY",
			Opcode::InverseQuery => "IQUERY",
			Opcode::ServerStatusRequest => "STATUS",
			_ => "UNKNOWN",
		};

		let mut flags: Vec<&str> = Vec::new();
		if p.header.authoritative {
			flags.push("AA");
		}
		if p.header.truncated {
			flags.push("TC");
		}
		if p.header.recursion_desired {
			flags.push("RD");
		}
		if p.header.recursion_available {
			flags.push("RA");
		}
		if p.header.authenticated_data {
			flags.push("AD");
		}
		if p.header.checking_disabled {
			flags.push("CD");
		}

		DnsHeader {
			id: Number::from(p.header.id),
			opcode: opcode_string.to_string(),
			flags: Value::from(flags),
			total_questions: Number::from(p.header.questions),
			total_answer_rrs: Number::from(p.header.answers),
			total_authority_rrs: Number::from(p.header.nameservers),
			total_additional_rrs: Number::from(p.header.additional),
		}
	}

	pub fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("id".to_string(), Value::Number(self.id.clone()));
		header.insert("opcode".to_string(), Value::String(self.opcode.clone()));
		header.insert("flags".to_string(), self.flags.clone());
		header.insert("total_questions".to_string(), Value::Number(self.total_questions.clone()));
		header.insert("total_answer_rrs".to_string(), Value::Number(self.total_answer_rrs.clone()));
		header.insert("total_authority_rrs".to_string(), Value::Number(self.total_authority_rrs.clone()));
		header.insert("total_additional_rrs".to_string(), Value::Number(self.total_additional_rrs.clone()));

		header
	}
}<commit_msg>Parse DNS Rcode<commit_after>use dns_parser::Packet as DnsPacket;
use dns_parser::{Opcode, ResponseCode};
use serde_json::{Value, Map, Number};

pub struct DnsHeader {
	id: Number,
	opcode: String,
	flags: Value,
	rcode: String,
	total_questions: Number,
	total_answer_rrs: Number,
	total_authority_rrs: Number,
	total_additional_rrs: Number,
}

impl DnsHeader {
	pub fn new(p: &DnsPacket) -> DnsHeader {
		let opcode_number = p.header.opcode;
		let opcode_string = match opcode_number {
			Opcode::StandardQuery => "QUERY",
			Opcode::InverseQuery => "IQUERY",
			Opcode::ServerStatusRequest => "STATUS",
			_ => "UNKNOWN",
		};

		let mut flags: Vec<&str> = Vec::new();
		if p.header.authoritative {
			flags.push("AA");
		}
		if p.header.truncated {
			flags.push("TC");
		}
		if p.header.recursion_desired {
			flags.push("RD");
		}
		if p.header.recursion_available {
			flags.push("RA");
		}
		if p.header.authenticated_data {
			flags.push("AD");
		}
		if p.header.checking_disabled {
			flags.push("CD");
		}

		let rcode_number = p.header.response_code;
		let rcode_string = match rcode_number {
			ResponseCode::NoError => "NOERROR",
			ResponseCode::FormatError => "FORMERR",
			ResponseCode::ServerFailure => "SERVFAIL",
			ResponseCode::NameError => "NXDOMAIN",
			ResponseCode::NotImplemented => "NOTIMPL",
			ResponseCode::Refused => "REFUSED",
			_ => "UNKNOWN",
		};

		DnsHeader {
			id: Number::from(p.header.id),
			opcode: opcode_string.to_string(),
			flags: Value::from(flags),
			rcode: rcode_string.to_string(),
			total_questions: Number::from(p.header.questions),
			total_answer_rrs: Number::from(p.header.answers),
			total_authority_rrs: Number::from(p.header.nameservers),
			total_additional_rrs: Number::from(p.header.additional),
		}
	}

	pub fn to_json_map(&self) -> Map<String, Value> {
		let mut header = Map::new();

		header.insert("id".to_string(), Value::Number(self.id.clone()));
		header.insert("opcode".to_string(), Value::String(self.opcode.clone()));
		header.insert("flags".to_string(), self.flags.clone());
		header.insert("rcode".to_string(), Value::String(self.rcode.clone()));
		header.insert("total_questions".to_string(), Value::Number(self.total_questions.clone()));
		header.insert("total_answer_rrs".to_string(), Value::Number(self.total_answer_rrs.clone()));
		header.insert("total_authority_rrs".to_string(), Value::Number(self.total_authority_rrs.clone()));
		header.insert("total_additional_rrs".to_string(), Value::Number(self.total_additional_rrs.clone()));

		header
	}
}<|endoftext|>