<commit_before>//! Virtualize processor hardware for multiple software environments using Virtual Machine Extensions.

use crate::bits64::rflags::{self, RFlags};
use crate::vmx::{Result, VmFail};
use core::arch::asm;

/// Helper used to extract VMX-specific Result in accordance with
/// conventions described in Intel SDM, Volume 3C, Section 30.2.
// We inline this to provide an obstruction-free path from this function's
// call site to the moment where `rflags::read()` reads RFLAGS. Otherwise it's
// possible for RFLAGS register to be clobbered by a function prologue,
// see https://github.com/gz/rust-x86/pull/50.
#[inline(always)]
fn vmx_capture_status() -> Result<()> {
    let flags = rflags::read();

    if flags.contains(RFlags::FLAGS_ZF) {
        Err(VmFail::VmFailValid)
    } else if flags.contains(RFlags::FLAGS_CF) {
        Err(VmFail::VmFailInvalid)
    } else {
        Ok(())
    }
}

/// Enable VMX operation.
///
/// `addr` specifies a 4KB-aligned physical address of VMXON region initialized
/// in accordance with Intel SDM, Volume 3C, Section 24.11.5.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmxon(addr: u64) -> Result<()> {
    asm!("vmxon ({0})", in(reg) addr, options(att_syntax));
    vmx_capture_status()
}

/// Disable VMX operation.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmxoff() -> Result<()> {
    asm!("vmxoff");
    vmx_capture_status()
}

/// Clear VMCS.
///
/// Ensures that VMCS data maintained on the processor is copied to the VMCS region
/// located at 4KB-aligned physical address `addr` and initializes some parts of it.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmclear(addr: u64) -> Result<()> {
    asm!("vmclear ({0})", in(reg) addr, options(att_syntax));
    vmx_capture_status()
}

/// Load current VMCS pointer.
///
/// Marks the current-VMCS pointer valid and loads it with the physical address `addr`.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmptrld(addr: u64) -> Result<()> {
    asm!("vmptrld ({0})", in(reg) addr, options(att_syntax));
    vmx_capture_status()
}

/// Return current VMCS pointer.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmptrst() -> Result<u64> {
    let value: u64 = 0;
    asm!("vmptrst ({0})", in(reg) &value, options(att_syntax));
    vmx_capture_status().and(Ok(value))
}

/// Read a specified field from a VMCS.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmread(field: u32) -> Result<u64> {
    let field: u64 = field.into();
    let value: u64;
    asm!("vmread {1}, {0}", in(reg) field, out(reg) value, options(att_syntax));
    vmx_capture_status().and(Ok(value))
}

/// Write to a specified field in a VMCS.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmwrite(field: u32, value: u64) -> Result<()> {
    let field: u64 = field.into();
    asm!("vmwrite {1}, {0}", in(reg) field, in(reg) value, options(att_syntax));
    vmx_capture_status()
}

/// Launch virtual machine.
///
/// # Safety
/// Needs CPL 0.
#[inline(always)]
pub unsafe fn vmlaunch() -> Result<()> {
    asm!("vmlaunch");
    vmx_capture_status()
}

/// Resume virtual machine.
///
/// # Safety
/// Needs CPL 0.
#[inline(always)]
pub unsafe fn vmresume() -> Result<()> {
    asm!("vmresume");
    vmx_capture_status()
}
<commit_msg>Fix asm syntax of some VMX instructions<commit_after>//! Virtualize processor hardware for multiple software environments using Virtual Machine Extensions.

use crate::bits64::rflags::{self, RFlags};
use crate::vmx::{Result, VmFail};
use core::arch::asm;

/// Helper used to extract VMX-specific Result in accordance with
/// conventions described in Intel SDM, Volume 3C, Section 30.2.
// We inline this to provide an obstruction-free path from this function's
// call site to the moment where `rflags::read()` reads RFLAGS. Otherwise it's
// possible for RFLAGS register to be clobbered by a function prologue,
// see https://github.com/gz/rust-x86/pull/50.
#[inline(always)]
fn vmx_capture_status() -> Result<()> {
    let flags = rflags::read();

    if flags.contains(RFlags::FLAGS_ZF) {
        Err(VmFail::VmFailValid)
    } else if flags.contains(RFlags::FLAGS_CF) {
        Err(VmFail::VmFailInvalid)
    } else {
        Ok(())
    }
}

/// Enable VMX operation.
///
/// `addr` specifies a 4KB-aligned physical address of VMXON region initialized
/// in accordance with Intel SDM, Volume 3C, Section 24.11.5.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmxon(addr: u64) -> Result<()> {
    asm!("vmxon ({0})", in(reg) &addr, options(att_syntax));
    vmx_capture_status()
}

/// Disable VMX operation.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmxoff() -> Result<()> {
    asm!("vmxoff");
    vmx_capture_status()
}

/// Clear VMCS.
///
/// Ensures that VMCS data maintained on the processor is copied to the VMCS region
/// located at 4KB-aligned physical address `addr` and initializes some parts of it.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmclear(addr: u64) -> Result<()> {
    asm!("vmclear ({0})", in(reg) &addr, options(att_syntax));
    vmx_capture_status()
}

/// Load current VMCS pointer.
///
/// Marks the current-VMCS pointer valid and loads it with the physical address `addr`.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmptrld(addr: u64) -> Result<()> {
    asm!("vmptrld ({0})", in(reg) &addr, options(att_syntax));
    vmx_capture_status()
}

/// Return current VMCS pointer.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmptrst() -> Result<u64> {
    let value: u64 = 0;
    asm!("vmptrst ({0})", in(reg) &value, options(att_syntax));
    vmx_capture_status().and(Ok(value))
}

/// Read a specified field from a VMCS.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmread(field: u32) -> Result<u64> {
    let field: u64 = field.into();
    let value: u64;
    asm!("vmread {0}, {1}", in(reg) field, out(reg) value, options(att_syntax));
    vmx_capture_status().and(Ok(value))
}

/// Write to a specified field in a VMCS.
///
/// # Safety
/// Needs CPL 0.
pub unsafe fn vmwrite(field: u32, value: u64) -> Result<()> {
    let field: u64 = field.into();
    asm!("vmwrite {1}, {0}", in(reg) field, in(reg) value, options(att_syntax));
    vmx_capture_status()
}

/// Launch virtual machine.
///
/// # Safety
/// Needs CPL 0.
#[inline(always)]
pub unsafe fn vmlaunch() -> Result<()> {
    asm!("vmlaunch");
    vmx_capture_status()
}

/// Resume virtual machine.
///
/// # Safety
/// Needs CPL 0.
#[inline(always)]
pub unsafe fn vmresume() -> Result<()> {
    asm!("vmresume");
    vmx_capture_status()
}
<|endoftext|>


<commit_before>extern crate glob;

use std::env;
use std::path::{Path};
use std::process::{Command};

// Environment variables:
//
// * LLVM_CONFIG_PATH - provides a path to an `llvm-config` executable
// * LIBCLANG_PATH - provides a path to a directory containing a `libclang` shared library
// * LIBCLANG_STATIC_PATH - provides a path to a directory containing LLVM and Clang static libraries

/// Returns whether the supplied directory contains the supplied file.
fn contains(directory: &str, file: &str) -> bool {
    Path::new(directory).join(file).exists()
}

/// Panics with a user friendly error message.
fn error(file: &str, env: &str) -> ! {
    panic!("could not find {0}, set the {1} environment variable to a path where {0} can be found", file, env);
}

/// Runs a console command, returning the output if the command was successfully executed.
fn run(command: &str, arguments: &[&str]) -> Option<String> {
    Command::new(command).args(arguments).output().map(|o| {
        String::from_utf8_lossy(&o.stdout).into_owned()
    }).ok()
}

/// Runs `llvm-config`, returning the output if the command was successfully executed.
fn run_llvm_config(arguments: &[&str]) -> Option<String> {
    run(&env::var("LLVM_CONFIG_PATH").unwrap_or("llvm-config".into()), arguments)
}

/// Backup search directories for FreeBSD and Linux.
const SEARCH_LINUX: &'static [&'static str] = &[
    "/usr/lib",
    "/usr/lib/llvm",
    "/usr/lib/llvm-3.8/lib",
    "/usr/lib/llvm-3.7/lib",
    "/usr/lib/llvm-3.6/lib",
    "/usr/lib/llvm-3.5/lib",
    "/usr/lib/llvm-3.4/lib",
    "/usr/local/llvm38/lib",
    "/usr/local/llvm37/lib",
    "/usr/local/llvm36/lib",
    "/usr/local/llvm35/lib",
    "/usr/local/llvm34/lib",
    "/usr/lib64/llvm",
    "/usr/lib/x86_64-linux-gnu",
    "/usr/local/lib",
];

/// Backup search directories for OS X.
const SEARCH_OSX: &'static [&'static str] = &[
    "/usr/local/opt/llvm/lib",
    "/Library/Developer/CommandLineTools/usr/lib",
    "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib",
    "/usr/local/opt/llvm35/lib/llvm-3.8/lib",
    "/usr/local/opt/llvm35/lib/llvm-3.7/lib",
    "/usr/local/opt/llvm35/lib/llvm-3.6/lib",
    "/usr/local/opt/llvm35/lib/llvm-3.5/lib",
];

/// Backup search directories for Windows.
const SEARCH_WINDOWS: &'static [&'static str] = &[
    "C:\\Program Files\\LLVM\\bin",
    "C:\\Program Files\\LLVM\\lib",
];

/// Searches for a library, returning the directory it can be found in if the search was successful.
fn find(file: &str, env: &str) -> Option<String> {
    // Search the directory provided by the relevant environment variable, if set.
    if let Some(directory) = env::var(env).ok() {
        if contains(&directory, file) {
            return Some(directory);
        }
    }

    // Search the `bin` and `lib` subdirectories in the directory returned by
    // `llvm-config --prefix`, if `llvm-config` is available.
    if let Some(output) = run_llvm_config(&["--prefix"]) {
        let directory = output.lines().map(|s| s.to_string()).next().unwrap();
        let bin = format!("{}/bin", directory);
        if contains(&bin, file) {
            return Some(bin);
        }
        let lib = format!("{}/lib", directory);
        if contains(&lib, file) {
            return Some(lib);
        }
    }

    // Search the backup directories.
    let search = if cfg!(any(target_os="freebsd", target_os="linux")) {
        SEARCH_LINUX
    } else if cfg!(target_os="macos") {
        SEARCH_OSX
    } else if cfg!(target_os="windows") {
        SEARCH_WINDOWS
    } else {
        return None;
    };
    search.iter().find(|d| contains(d, file)).map(|s| s.to_string())
}

/// Returns the name of an LLVM or Clang library from a path.
fn get_library_name(path: &Path) -> Option<String> {
    path.file_stem().map(|l| l.to_string_lossy()[3..].into())
}

/// Returns the LLVM libraries required to link to `libclang` statically.
fn get_llvm_libraries() -> Vec<String> {
    run_llvm_config(&["--libs"]).expect(
        "could not execute `llvm-config --libs`, set the LLVM_CONFIG_PATH environment variable to \
         a path to an `llvm-config` executable"
    ).split_whitespace().filter_map(|p| {
        // Depending on the version of `llvm-config` in use, listed libraries may be in one of two
        // forms, a full path to the library or simply prefixed with `-l`.
        if p.starts_with("-l") {
            Some(p[2..].into())
        } else {
            get_library_name(&Path::new(p))
        }
    }).collect()
}

/// Clang libraries required to link to `libclang` 3.5 and later statically.
const CLANG_LIBRARIES: &'static [&'static str] = &[
    "clang",
    "clangAST",
    "clangAnalysis",
    "clangBasic",
    "clangDriver",
    "clangEdit",
    "clangFrontend",
    "clangIndex",
    "clangLex",
    "clangParse",
    "clangRewrite",
    "clangSema",
    "clangSerialization",
];

/// Returns the Clang libraries required to link to `libclang` statically.
fn get_clang_libraries(directory: &str) -> Vec<String> {
    let pattern = Path::new(directory).join("libclang*.a").to_string_lossy().to_string();
    if let Ok(libraries) = glob::glob(&pattern) {
        libraries.filter_map(|l| l.ok().and_then(|l| get_library_name(&l))).collect()
    } else {
        CLANG_LIBRARIES.iter().map(|l| l.to_string()).collect()
    }
}

fn main() {
    if cfg!(feature="static") {
        // Find LLVM and Clang static libraries.
        let directory = match find("libclang.a", "LIBCLANG_STATIC_PATH") {
            Some(directory) => directory,
            _ => error("libclang.a", "LIBCLANG_STATIC_PATH"),
        };

        print!("cargo:rustc-flags=");

        // Specify required LLVM and Clang static libraries.
        print!("-L {} ", directory);
        for library in get_llvm_libraries() {
            print!("-l static={} ", library);
        }
        for library in get_clang_libraries(&directory) {
            print!("-l static={} ", library);
        }

        // Specify required system libraries.
        if cfg!(any(target_os="freebsd", target_os="linux")) {
            println!("-l ffi -l ncursesw -l stdc++ -l z");
        } else if cfg!(target_os="macos") {
            println!("-l ffi -l ncurses -l stdc++ -l z");
        } else {
            panic!("unsupported operating system for static linking");
        }
    } else {
        let file = if cfg!(target_os="windows") {
            // The filename of the `libclang` shared library on Windows is `libclang.dll` instead of
            // the expected `clang.dll`.
            "libclang.dll".into()
        } else {
            format!("{}clang{}", env::consts::DLL_PREFIX, env::consts::DLL_SUFFIX)
        };

        // Find the `libclang` shared library.
        let directory = match find(&file, "LIBCLANG_PATH") {
            Some(directory) => directory,
            _ => error(&file, "LIBCLANG_PATH")
        };

        println!("cargo:rustc-link-search={}", directory);
        if cfg!(all(target_os="windows", target_env="msvc")) {
            // Find the `libclang` stub static library required for the MSVC toolchain.
            let directory = match find("libclang.lib", "LIBCLANG_PATH") {
                Some(directory) => directory,
                _ => error("libclang.lib", "LIBCLANG_PATH"),
            };

            println!("cargo:rustc-link-search={}", directory);
            println!("cargo:rustc-link-lib=dylib=libclang");
        } else {
            println!("cargo:rustc-link-lib=dylib=clang");
        }
    }
}
<commit_msg>Use globs to define backup search directories<commit_after>extern crate glob;

use std::env;
use std::path::{Path};
use std::process::{Command};

use glob::{MatchOptions};

// Environment variables:
//
// * LLVM_CONFIG_PATH - provides a path to an `llvm-config` executable
// * LIBCLANG_PATH - provides a path to a directory containing a `libclang` shared library
// * LIBCLANG_STATIC_PATH - provides a path to a directory containing LLVM and Clang static libraries

/// Returns whether the supplied directory contains the supplied file.
fn contains<D: AsRef<Path>>(directory: D, file: &str) -> bool {
    directory.as_ref().join(file).exists()
}

/// Panics with a user friendly error message.
fn error(file: &str, env: &str) -> ! {
    panic!("could not find {0}, set the {1} environment variable to a path where {0} can be found", file, env);
}

/// Runs a console command, returning the output if the command was successfully executed.
fn run(command: &str, arguments: &[&str]) -> Option<String> {
    Command::new(command).args(arguments).output().map(|o| {
        String::from_utf8_lossy(&o.stdout).into_owned()
    }).ok()
}

/// Runs `llvm-config`, returning the output if the command was successfully executed.
fn run_llvm_config(arguments: &[&str]) -> Option<String> {
    run(&env::var("LLVM_CONFIG_PATH").unwrap_or("llvm-config".into()), arguments)
}

/// Backup search directory globs for FreeBSD and Linux.
const SEARCH_LINUX: &'static [&'static str] = &[
    "/usr/*",
    "/usr/*/*",
    "/usr/*/*/*",
];

/// Backup search directory globs for OS X.
const SEARCH_OSX: &'static [&'static str] = &[
    "/usr/local/opt/*/lib",
    "/Library/Developer/CommandLineTools/usr/lib",
    "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib",
    "/usr/local/opt/*/lib/*/lib",
];

/// Backup search directory globs for Windows.
const SEARCH_WINDOWS: &'static [&'static str] = &[
    "C:\\LLVM\\bin",
    "C:\\LLVM\\lib",
    "C:\\Program Files*\\LLVM\\bin",
    "C:\\Program Files*\\LLVM\\lib",
];

/// Searches for a library, returning the directory it can be found in if the search was successful.
fn find(file: &str, env: &str) -> Option<String> {
    // Search the directory provided by the relevant environment variable, if set.
    if let Some(directory) = env::var(env).ok() {
        if contains(&directory, file) {
            return Some(directory);
        }
    }

    // Search the `bin` and `lib` subdirectories in the directory returned by
    // `llvm-config --prefix`, if `llvm-config` is available.
    if let Some(output) = run_llvm_config(&["--prefix"]) {
        let directory = output.lines().map(|s| s.to_string()).next().unwrap();
        let bin = format!("{}/bin", directory);
        if contains(&bin, file) {
            return Some(bin);
        }
        let lib = format!("{}/lib", directory);
        if contains(&lib, file) {
            return Some(lib);
        }
    }

    // Search the backup directories.
    let search = if cfg!(any(target_os="freebsd", target_os="linux")) {
        SEARCH_LINUX
    } else if cfg!(target_os="macos") {
        SEARCH_OSX
    } else if cfg!(target_os="windows") {
        SEARCH_WINDOWS
    } else {
        return None;
    };
    for pattern in search {
        let mut options = MatchOptions::new();
        options.case_sensitive = false;
        options.require_literal_separator = true;
        if let Ok(paths) = glob::glob_with(pattern, &options) {
            for path in paths.filter_map(Result::ok) {
                if contains(&path, file) {
                    return Some(path.to_string_lossy().into_owned());
                }
            }
        }
    }
    None
}

/// Returns the name of an LLVM or Clang library from a path.
fn get_library_name(path: &Path) -> Option<String> {
    path.file_stem().map(|l| l.to_string_lossy()[3..].into())
}

/// Returns the LLVM libraries required to link to `libclang` statically.
fn get_llvm_libraries() -> Vec<String> {
    run_llvm_config(&["--libs"]).expect(
        "could not execute `llvm-config --libs`, set the LLVM_CONFIG_PATH environment variable to \
         a path to an `llvm-config` executable"
    ).split_whitespace().filter_map(|p| {
        // Depending on the version of `llvm-config` in use, listed libraries may be in one of two
        // forms, a full path to the library or simply prefixed with `-l`.
        if p.starts_with("-l") {
            Some(p[2..].into())
        } else {
            get_library_name(&Path::new(p))
        }
    }).collect()
}

/// Clang libraries required to link to `libclang` 3.5 and later statically.
const CLANG_LIBRARIES: &'static [&'static str] = &[
    "clang",
    "clangAST",
    "clangAnalysis",
    "clangBasic",
    "clangDriver",
    "clangEdit",
    "clangFrontend",
    "clangIndex",
    "clangLex",
    "clangParse",
    "clangRewrite",
    "clangSema",
    "clangSerialization",
];

/// Returns the Clang libraries required to link to `libclang` statically.
fn get_clang_libraries(directory: &str) -> Vec<String> {
    let pattern = Path::new(directory).join("libclang*.a").to_string_lossy().to_string();
    if let Ok(libraries) = glob::glob(&pattern) {
        libraries.filter_map(|l| l.ok().and_then(|l| get_library_name(&l))).collect()
    } else {
        CLANG_LIBRARIES.iter().map(|l| l.to_string()).collect()
    }
}

fn main() {
    if cfg!(feature="static") {
        // Find LLVM and Clang static libraries.
        let directory = match find("libclang.a", "LIBCLANG_STATIC_PATH") {
            Some(directory) => directory,
            _ => error("libclang.a", "LIBCLANG_STATIC_PATH"),
        };

        print!("cargo:rustc-flags=");

        // Specify required LLVM and Clang static libraries.
        print!("-L {} ", directory);
        for library in get_llvm_libraries() {
            print!("-l static={} ", library);
        }
        for library in get_clang_libraries(&directory) {
            print!("-l static={} ", library);
        }

        // Specify required system libraries.
        if cfg!(any(target_os="freebsd", target_os="linux")) {
            println!("-l ffi -l ncursesw -l stdc++ -l z");
        } else if cfg!(target_os="macos") {
            println!("-l ffi -l ncurses -l stdc++ -l z");
        } else {
            panic!("unsupported operating system for static linking");
        }
    } else {
        let file = if cfg!(target_os="windows") {
            // The filename of the `libclang` shared library on Windows is `libclang.dll` instead of
            // the expected `clang.dll`.
            "libclang.dll".into()
        } else {
            format!("{}clang{}", env::consts::DLL_PREFIX, env::consts::DLL_SUFFIX)
        };

        // Find the `libclang` shared library.
        let directory = match find(&file, "LIBCLANG_PATH") {
            Some(directory) => directory,
            _ => error(&file, "LIBCLANG_PATH")
        };

        println!("cargo:rustc-link-search={}", directory);
        if cfg!(all(target_os="windows", target_env="msvc")) {
            // Find the `libclang` stub static library required for the MSVC toolchain.
            let directory = match find("libclang.lib", "LIBCLANG_PATH") {
                Some(directory) => directory,
                _ => error("libclang.lib", "LIBCLANG_PATH"),
            };

            println!("cargo:rustc-link-search={}", directory);
            println!("cargo:rustc-link-lib=dylib=libclang");
        } else {
            println!("cargo:rustc-link-lib=dylib=clang");
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2021, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

#![allow(clippy::print_literal)]
#![allow(clippy::unused_io_amount)]

use rustc_version::{version, version_meta, Channel, Version};
#[allow(unused_imports)]
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::exit;

#[allow(dead_code)]
fn rerun_dir<P: AsRef<Path>>(dir: P) {
  for entry in fs::read_dir(dir).unwrap() {
    let entry = entry.unwrap();
    let path = entry.path();
    println!("cargo:rerun-if-changed={}", path.to_string_lossy());

    if path.is_dir() {
      rerun_dir(path);
    }
  }
}

#[allow(dead_code)]
fn hash_changed(
  files: &[&str], out_dir: &str, config: &Path,
) -> Option<([u8; 8], PathBuf)> {
  use std::collections::hash_map::DefaultHasher;
  use std::hash::Hasher;
  use std::io::Read;

  let mut hasher = DefaultHasher::new();

  let paths = files
    .iter()
    .map(Path::new)
    .chain(std::iter::once(config))
    .chain(std::iter::once(Path::new("build.rs")));

  for path in paths {
    if let Ok(mut f) = std::fs::File::open(path) {
      let mut buf = Vec::new();
      f.read_to_end(&mut buf).unwrap();

      hasher.write(&buf);
    } else {
      panic!("Cannot open {}", path.display());
    }
  }

  let strip = env::var("STRIP").unwrap_or_else(|_| "strip".to_string());

  hasher.write(strip.as_bytes());

  let hash = hasher.finish().to_be_bytes();

  let hash_path = Path::new(&out_dir).join("asm.hash");

  if let Ok(old_hash) = std::fs::read(&hash_path) {
    if old_hash == hash {
      return None;
    }
  }

  Some((hash, hash_path))
}

#[cfg(feature = "asm")]
fn build_nasm_files() {
  use std::fs::File;
  use std::io::Write;
  let out_dir = env::var("OUT_DIR").unwrap();

  let dest_path = Path::new(&out_dir).join("config.asm");
  let mut config_file = File::create(&dest_path).unwrap();
  config_file.write(b"	%define private_prefix rav1e\n").unwrap();
  config_file.write(b"	%define ARCH_X86_32 0\n").unwrap();
  config_file.write(b" %define ARCH_X86_64 1\n").unwrap();
  config_file.write(b"	%define PIC 1\n").unwrap();
  config_file.write(b" %define STACK_ALIGNMENT 16\n").unwrap();
  config_file.write(b" %define HAVE_AVX512ICL 1\n").unwrap();
  if env::var("CARGO_CFG_TARGET_VENDOR").unwrap() == "apple" {
    config_file.write(b" %define PREFIX 1\n").unwrap();
  }

  let asm_files = &[
    "src/x86/ipred_avx2.asm",
    "src/x86/ipred_sse.asm",
    "src/x86/ipred16_avx2.asm",
    "src/x86/itx_avx2.asm",
    "src/x86/itx_sse.asm",
    "src/x86/itx16_avx2.asm",
    "src/x86/itx16_sse.asm",
    "src/x86/looprestoration_avx2.asm",
    "src/x86/looprestoration16_avx2.asm",
    "src/x86/mc_avx2.asm",
    "src/x86/mc16_avx2.asm",
    "src/x86/mc_avx512.asm",
    "src/x86/mc_sse.asm",
    "src/x86/mc16_sse.asm",
    "src/x86/me.asm",
    "src/x86/sad_sse2.asm",
    "src/x86/sad_avx.asm",
    "src/x86/satd.asm",
    "src/x86/sse.asm",
    "src/x86/cdef_rav1e.asm",
    "src/x86/cdef_sse.asm",
    "src/x86/cdef16_avx2.asm",
    "src/x86/cdef16_sse.asm",
    "src/x86/tables.asm",
  ];

  if let Some((hash, hash_path)) =
    hash_changed(asm_files, &out_dir, &dest_path)
  {
    let mut config_include_arg = String::from("-I");
    config_include_arg.push_str(&out_dir);
    config_include_arg.push('/');
    let mut nasm = nasm_rs::Build::new();
    for file in asm_files {
      nasm.file(file);
    }
    nasm.flag(&config_include_arg);
    nasm.flag("-Isrc/");
    let obj = nasm.compile_objects().expect(
      "NASM build failed. Make sure you have nasm installed. https://nasm.us",
    );
    // cc is better at finding the correct archiver
    let mut cc = cc::Build::new();
    for o in obj {
      cc.object(o);
    }
    cc.compile("rav1easm");

    // Strip local symbols from the asm library since they
    // confuse the debugger.
    fn strip<P: AsRef<Path>>(obj: P) {
      let strip = env::var("STRIP").unwrap_or_else(|_| "strip".to_string());

      let mut cmd = std::process::Command::new(strip);

      cmd.arg("-x").arg(obj.as_ref());

      let _ = cmd.output();
    }

    strip(Path::new(&out_dir).join("librav1easm.a"));

    std::fs::write(hash_path, &hash[..]).unwrap();
  } else {
    println!("cargo:rustc-link-search={}", out_dir);
  }
  println!("cargo:rustc-link-lib=static=rav1easm");
  rerun_dir("src/x86");
  rerun_dir("src/ext/x86");
}

#[cfg(feature = "asm")]
fn build_asm_files() {
  use std::fs::File;
  use std::io::Write;
  let out_dir = env::var("OUT_DIR").unwrap();

  let dest_path = Path::new(&out_dir).join("config.h");
  let mut config_file = File::create(&dest_path).unwrap();
  if env::var("CARGO_CFG_TARGET_VENDOR").unwrap() == "apple" {
    config_file.write(b" #define PREFIX 1\n").unwrap();
  }
  config_file.write(b" #define PRIVATE_PREFIX rav1e_\n").unwrap();
  config_file.write(b" #define ARCH_AARCH64 1\n").unwrap();
  config_file.write(b" #define ARCH_ARM 0\n").unwrap();
  config_file.write(b" #define CONFIG_LOG 1 \n").unwrap();
  config_file.write(b" #define HAVE_ASM 1\n").unwrap();
  config_file.sync_all().unwrap();

  let asm_files = &[
    "src/arm/64/cdef.S",
    "src/arm/64/cdef16.S",
    "src/arm/64/mc.S",
    "src/arm/64/mc16.S",
    "src/arm/64/itx.S",
    "src/arm/64/itx16.S",
    "src/arm/64/ipred.S",
    "src/arm/64/ipred16.S",
    "src/arm/64/sad.S",
    "src/arm/tables.S",
  ];

  if let Some((hash, hash_path)) =
    hash_changed(asm_files, &out_dir, &dest_path)
  {
    cc::Build::new()
      .files(asm_files)
      .include(".")
      .include(&out_dir)
      .compile("rav1e-aarch64");

    std::fs::write(hash_path, &hash[..]).unwrap();
  } else {
    println!("cargo:rustc-link-search={}", out_dir);
    println!("cargo:rustc-link-lib=static=rav1e-aarch64");
  }
  rerun_dir("src/arm");
}

fn rustc_version_check() {
  // This should match the version in the CI
  // Make sure to updated README.md when this changes.
  const REQUIRED_VERSION: &str = "1.51.0";
  if version().unwrap() < Version::parse(REQUIRED_VERSION).unwrap() {
    eprintln!("rav1e requires rustc >= {}.", REQUIRED_VERSION);
    exit(1);
  }

  if version_meta().unwrap().channel == Channel::Nightly {
    println!("cargo:rustc-cfg=nightly_rustc");
  }
}

#[cfg(feature = "asm")]
fn is_nasm_new_enough(
  min_version: Option<(u8, u8, u8)>, nasm_path: &Path,
) -> Result<(), String> {
  match get_output(std::process::Command::new(nasm_path).arg("-v")) {
    Ok(version) => {
      if version.contains("NASM version 0.") {
        Err(version)
      } else if let Some((major, minor, micro)) = min_version {
        let ver = parse_nasm_version(&version)?;
        eprintln!("{:?} {:?} {:?}", min_version, version, ver);
        if major > ver.0
          || (major == ver.0 && minor > ver.1)
          || (major == ver.0 && minor == ver.1 && micro > ver.2)
        {
          Err(version)
        } else {
          Ok(())
        }
      } else {
        Ok(())
      }
    }
    Err(err) => Err(err),
  }
}

#[cfg(feature = "asm")]
fn parse_nasm_version(version_string: &str) -> Result<(u8, u8, u8), String> {
  // Rust regex lib doesn't support empty expressions in alterations,
  // so workaround the other way
  let regex1 =
    regex::Regex::new(r"(?:NASM version )?(\d+)\.(\d+)\.(\d+)").unwrap();
  if let Some(captures) = regex1.captures(version_string) {
    Ok((
      captures[1].parse().expect("Invalid version component"),
      captures[2].parse().expect("Invalid version component"),
      captures[3].parse().expect("Invalid version component"),
    ))
  } else {
    let regex2 = regex::Regex::new(r"(?:NASM version )?(\d+)\.(\d+)").unwrap();
    let captures = regex2
      .captures(version_string)
      .ok_or_else(|| "Unable to parse NASM version string".to_string())?;
    Ok((
      captures[1].parse().expect("Invalid version component"),
      captures[2].parse().expect("Invalid version component"),
      0,
    ))
  }
}

#[test]
fn test_parse_nasm_version() {
  let ver_str = "NASM version 2.14.02 compiled on Jan 22 2019";
  assert_eq!((2, 14, 2), parse_nasm_version(ver_str));
  let ver_str = "NASM version 2.14 compiled on Jan 22 2019";
  assert_eq!((2, 14, 0), parse_nasm_version(ver_str));
}

#[cfg(feature = "asm")]
fn is_nasm_installed_check() {
  let mut cmd = std::process::Command::new("nasm");
  cmd.arg("-h");
  match cmd.status() {
    Ok(_) => (),
    Err(_) => {
      panic!(
        r#"Couldn't find NASM, make sure NASM installed and in PATH or disable "asm" feature"#
      )
    }
  }
}

#[cfg(feature = "asm")]
fn nasm_version_check() -> std::path::PathBuf {
  let nasm_path = std::path::PathBuf::from("nasm");
  match is_nasm_new_enough(Some((2, 14, 0)), &nasm_path) {
    Ok(_) => nasm_path,
    Err(version) => {
      panic!("This version of NASM is too old: {}", version);
    }
  }
}

#[cfg(feature = "asm")]
fn get_output(cmd: &mut std::process::Command) -> Result<String, String> {
  let out = cmd.output().map_err(|e| e.to_string())?;
  if out.status.success() {
    Ok(String::from_utf8_lossy(&out.stdout).to_string())
  } else {
    Err(String::from_utf8_lossy(&out.stderr).to_string())
  }
}

#[allow(unused_variables)]
fn main() {
  rustc_version_check();

  let arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
  let os = env::var("CARGO_CFG_TARGET_OS").unwrap();
  // let env = env::var("CARGO_CFG_TARGET_ENV").unwrap();

  #[cfg(feature = "asm")]
  {
    if arch == "x86_64" {
      println!("cargo:rustc-cfg={}", "nasm_x86_64");
      is_nasm_installed_check();
      nasm_version_check();
      build_nasm_files()
    }
    if arch == "aarch64" {
      println!("cargo:rustc-cfg={}", "asm_neon");
      build_asm_files()
    }
  }

  if os == "windows" && cfg!(feature = "decode_test") {
    panic!("Unsupported feature on this platform!");
  }

  vergen::generate_cargo_keys(vergen::ConstantsFlags::all())
    .expect("Unable to generate the cargo keys!");

  println!("cargo:rustc-env=PROFILE={}", env::var("PROFILE").unwrap());
}
<commit_msg>Use nasm-rs for nasm version check<commit_after>// Copyright (c) 2017-2021, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

#![allow(clippy::print_literal)]
#![allow(clippy::unused_io_amount)]

use rustc_version::{version, version_meta, Channel, Version};
#[allow(unused_imports)]
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::exit;

#[allow(dead_code)]
fn rerun_dir<P: AsRef<Path>>(dir: P) {
  for entry in fs::read_dir(dir).unwrap() {
    let entry = entry.unwrap();
    let path = entry.path();
    println!("cargo:rerun-if-changed={}", path.to_string_lossy());

    if path.is_dir() {
      rerun_dir(path);
    }
  }
}

#[allow(dead_code)]
fn hash_changed(
  files: &[&str], out_dir: &str, config: &Path,
) -> Option<([u8; 8], PathBuf)> {
  use std::collections::hash_map::DefaultHasher;
  use std::hash::Hasher;
  use std::io::Read;

  let mut hasher = DefaultHasher::new();

  let paths = files
    .iter()
    .map(Path::new)
    .chain(std::iter::once(config))
    .chain(std::iter::once(Path::new("build.rs")));

  for path in paths {
    if let Ok(mut f) = std::fs::File::open(path) {
      let mut buf = Vec::new();
      f.read_to_end(&mut buf).unwrap();

      hasher.write(&buf);
    } else {
      panic!("Cannot open {}", path.display());
    }
  }

  let strip = env::var("STRIP").unwrap_or_else(|_| "strip".to_string());

  hasher.write(strip.as_bytes());

  let hash = hasher.finish().to_be_bytes();

  let hash_path = Path::new(&out_dir).join("asm.hash");

  if let Ok(old_hash) = std::fs::read(&hash_path) {
    if old_hash == hash {
      return None;
    }
  }

  Some((hash, hash_path))
}

#[cfg(feature = "asm")]
fn build_nasm_files() {
  use std::fs::File;
  use std::io::Write;
  let out_dir = env::var("OUT_DIR").unwrap();

  let dest_path = Path::new(&out_dir).join("config.asm");
  let mut config_file = File::create(&dest_path).unwrap();
  config_file.write(b"	%define private_prefix rav1e\n").unwrap();
  config_file.write(b"	%define ARCH_X86_32 0\n").unwrap();
  config_file.write(b" %define ARCH_X86_64 1\n").unwrap();
  config_file.write(b"	%define PIC 1\n").unwrap();
  config_file.write(b" %define STACK_ALIGNMENT 16\n").unwrap();
  config_file.write(b" %define HAVE_AVX512ICL 1\n").unwrap();
  if env::var("CARGO_CFG_TARGET_VENDOR").unwrap() == "apple" {
    config_file.write(b" %define PREFIX 1\n").unwrap();
  }

  let asm_files = &[
    "src/x86/ipred_avx2.asm",
    "src/x86/ipred_sse.asm",
    "src/x86/ipred16_avx2.asm",
    "src/x86/itx_avx2.asm",
    "src/x86/itx_sse.asm",
    "src/x86/itx16_avx2.asm",
    "src/x86/itx16_sse.asm",
    "src/x86/looprestoration_avx2.asm",
    "src/x86/looprestoration16_avx2.asm",
    "src/x86/mc_avx2.asm",
    "src/x86/mc16_avx2.asm",
    "src/x86/mc_avx512.asm",
    "src/x86/mc_sse.asm",
    "src/x86/mc16_sse.asm",
    "src/x86/me.asm",
    "src/x86/sad_sse2.asm",
    "src/x86/sad_avx.asm",
    "src/x86/satd.asm",
    "src/x86/sse.asm",
    "src/x86/cdef_rav1e.asm",
    "src/x86/cdef_sse.asm",
    "src/x86/cdef16_avx2.asm",
    "src/x86/cdef16_sse.asm",
    "src/x86/tables.asm",
  ];

  if let Some((hash, hash_path)) =
    hash_changed(asm_files, &out_dir, &dest_path)
  {
    let mut config_include_arg = String::from("-I");
    config_include_arg.push_str(&out_dir);
    config_include_arg.push('/');
    let mut nasm = nasm_rs::Build::new();
    nasm.min_version(2, 14, 0);
    for file in asm_files {
      nasm.file(file);
    }
    nasm.flag(&config_include_arg);
    nasm.flag("-Isrc/");
    let obj = nasm.compile_objects().unwrap_or_else(|e| {
      println!("cargo:warning={}", e);
      panic!("NASM build failed. Make sure you have nasm installed or disable the \"asm\" feature.\n\
        You can get NASM from https://nasm.us or your system's package manager.\n\nerror: {}", e);
    });

    // cc is better at finding the correct archiver
    let mut cc = cc::Build::new();
    for o in obj {
      cc.object(o);
    }
    cc.compile("rav1easm");

    // Strip local symbols from the asm library since they
    // confuse the debugger.
    fn strip<P: AsRef<Path>>(obj: P) {
      let strip = env::var("STRIP").unwrap_or_else(|_| "strip".to_string());

      let mut cmd = std::process::Command::new(strip);

      cmd.arg("-x").arg(obj.as_ref());

      let _ = cmd.output();
    }

    strip(Path::new(&out_dir).join("librav1easm.a"));

    std::fs::write(hash_path, &hash[..]).unwrap();
  } else {
    println!("cargo:rustc-link-search={}", out_dir);
  }
  println!("cargo:rustc-link-lib=static=rav1easm");
  rerun_dir("src/x86");
  rerun_dir("src/ext/x86");
}

#[cfg(feature = "asm")]
fn build_asm_files() {
  use std::fs::File;
  use std::io::Write;
  let out_dir = env::var("OUT_DIR").unwrap();

  let dest_path = Path::new(&out_dir).join("config.h");
  let mut config_file = File::create(&dest_path).unwrap();
  if env::var("CARGO_CFG_TARGET_VENDOR").unwrap() == "apple" {
    config_file.write(b" #define PREFIX 1\n").unwrap();
  }
  config_file.write(b" #define PRIVATE_PREFIX rav1e_\n").unwrap();
  config_file.write(b" #define ARCH_AARCH64 1\n").unwrap();
  config_file.write(b" #define ARCH_ARM 0\n").unwrap();
  config_file.write(b" #define CONFIG_LOG 1 \n").unwrap();
  config_file.write(b" #define HAVE_ASM 1\n").unwrap();
  config_file.sync_all().unwrap();

  let asm_files = &[
    "src/arm/64/cdef.S",
    "src/arm/64/cdef16.S",
    "src/arm/64/mc.S",
    "src/arm/64/mc16.S",
    "src/arm/64/itx.S",
    "src/arm/64/itx16.S",
    "src/arm/64/ipred.S",
    "src/arm/64/ipred16.S",
    "src/arm/64/sad.S",
    "src/arm/tables.S",
  ];

  if let Some((hash, hash_path)) =
    hash_changed(asm_files, &out_dir, &dest_path)
  {
    cc::Build::new()
      .files(asm_files)
      .include(".")
      .include(&out_dir)
      .compile("rav1e-aarch64");

    std::fs::write(hash_path, &hash[..]).unwrap();
  } else {
    println!("cargo:rustc-link-search={}", out_dir);
    println!("cargo:rustc-link-lib=static=rav1e-aarch64");
  }
  rerun_dir("src/arm");
}

fn rustc_version_check() {
  // This should match the version in the CI
  // Make sure to updated README.md when this changes.
  const REQUIRED_VERSION: &str = "1.51.0";
  if version().unwrap() < Version::parse(REQUIRED_VERSION).unwrap() {
    eprintln!("rav1e requires rustc >= {}.", REQUIRED_VERSION);
    exit(1);
  }

  if version_meta().unwrap().channel == Channel::Nightly {
    println!("cargo:rustc-cfg=nightly_rustc");
  }
}

#[allow(unused_variables)]
fn main() {
  rustc_version_check();

  let arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
  let os = env::var("CARGO_CFG_TARGET_OS").unwrap();
  // let env = env::var("CARGO_CFG_TARGET_ENV").unwrap();

  #[cfg(feature = "asm")]
  {
    if arch == "x86_64" {
      println!("cargo:rustc-cfg={}", "nasm_x86_64");
      build_nasm_files()
    }
    if arch == "aarch64" {
      println!("cargo:rustc-cfg={}", "asm_neon");
      build_asm_files()
    }
  }

  if os == "windows" && cfg!(feature = "decode_test") {
    panic!("Unsupported feature on this platform!");
  }

  vergen::generate_cargo_keys(vergen::ConstantsFlags::all())
    .expect("Unable to generate the cargo keys!");

  println!("cargo:rustc-env=PROFILE={}", env::var("PROFILE").unwrap());
}
<|endoftext|>


<commit_before>extern crate csv;
extern crate phf_codegen;
extern crate rustc_serialize;

use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;

#[derive(RustcDecodable)]
struct TempRecord {
    zip_code: String,
    zip_code_type: String,
    city: String,
    state: String,
    _location_type: String,
    latitude: Option<f64>,
    longitude: Option<f64>,
    _location: String,
    is_decommissioned: String,
    tax_returns_filed: Option<u64>,
    estimated_population: Option<u64>,
    total_wages: Option<u64>,
}

impl TempRecord {
    fn to_syntax(&self) -> String {
        let mut output = String::from("Record {");
        output.push_str(&format!("zip_code: \"{}\",", self.zip_code));
        output.push_str(&format!("zip_code_type: {},", self.encoded_type()));
        output.push_str(&format!("city: \"{}\",", self.city));
        output.push_str(&format!("state: \"{}\",", self.state));
        output.push_str(&format!("coordinates: {},", self.coordinates()));
        output.push_str(&format!("is_decommissioned: {},", self.is_decommissioned));
        output.push_str(&format!("tax_returns_filed: {:?},", self.tax_returns_filed));
        output.push_str(&format!("estimated_population: {:?},", self.estimated_population));
        output.push_str(&format!("total_wages: {:?},", self.total_wages));
        output.push_str("}");
        output
    }

    fn encoded_type(&self) -> &'static str {
        match &self.zip_code_type[..] {
            "STANDARD" => "Type::Standard",
            "PO BOX" => "Type::PoBox",
            "UNIQUE" => "Type::Unique",
            "MILITARY" => "Type::Military",
            t => panic!("invalid ZIP code type \"{}\"", t),
        }
    }

    fn coordinates(&self) -> String {
        if let Some(ref latitude) = self.latitude {
            if let Some(ref longitude) = self.longitude {
                return format!("Some(({} as f64, {} as f64))", latitude, longitude);
            }
        }

        "None".into()
    }
}

fn main() {
    let path = Path::new(env!("OUT_DIR")).join("codegen.rs");
    let mut file = BufWriter::new(File::create(&path).unwrap());

    write!(&mut file, "static ZIP_CODES: phf::Map<&'static str, Record> = ").unwrap();
    let mut reader = csv::Reader::from_file("./data/zip_codes.csv").unwrap();
    let records = reader.decode().map(|r| r.unwrap()).collect::<Vec<TempRecord>>();

    {
        let mut map = phf_codegen::Map::new();
        for ref record in &records {
            map.entry(&record.zip_code[..], &record.to_syntax());
        }
        map.build(&mut file).unwrap();
    }

    write!(&mut file, ";\n").unwrap();
}
<commit_msg>Use env::var_os intead of env! in build script<commit_after>extern crate csv;
extern crate phf_codegen;
extern crate rustc_serialize;

use std::env;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;

#[derive(RustcDecodable)]
struct TempRecord {
    zip_code: String,
    zip_code_type: String,
    city: String,
    state: String,
    _location_type: String,
    latitude: Option<f64>,
    longitude: Option<f64>,
    _location: String,
    is_decommissioned: String,
    tax_returns_filed: Option<u64>,
    estimated_population: Option<u64>,
    total_wages: Option<u64>,
}

impl TempRecord {
    fn to_syntax(&self) -> String {
        let mut output = String::from("Record {");
        output.push_str(&format!("zip_code: \"{}\",", self.zip_code));
        output.push_str(&format!("zip_code_type: {},", self.encoded_type()));
        output.push_str(&format!("city: \"{}\",", self.city));
        output.push_str(&format!("state: \"{}\",", self.state));
        output.push_str(&format!("coordinates: {},", self.coordinates()));
        output.push_str(&format!("is_decommissioned: {},", self.is_decommissioned));
        output.push_str(&format!("tax_returns_filed: {:?},", self.tax_returns_filed));
        output.push_str(&format!("estimated_population: {:?},", self.estimated_population));
        output.push_str(&format!("total_wages: {:?},", self.total_wages));
        output.push_str("}");
        output
    }

    fn encoded_type(&self) -> &'static str {
        match &self.zip_code_type[..] {
            "STANDARD" => "Type::Standard",
            "PO BOX" => "Type::PoBox",
            "UNIQUE" => "Type::Unique",
            "MILITARY" => "Type::Military",
            t => panic!("invalid ZIP code type \"{}\"", t),
        }
    }

    fn coordinates(&self) -> String {
        if let Some(ref latitude) = self.latitude {
            if let Some(ref longitude) = self.longitude {
                return format!("Some(({} as f64, {} as f64))", latitude, longitude);
            }
        }

        "None".into()
    }
}

fn main() {
    let path = Path::new(&env::var_os("OUT_DIR").unwrap()).join("codegen.rs");
    let mut file = BufWriter::new(File::create(&path).unwrap());

    write!(&mut file, "static ZIP_CODES: phf::Map<&'static str, Record> = ").unwrap();
    let mut reader = csv::Reader::from_file("./data/zip_codes.csv").unwrap();
    let records = reader.decode().map(|r| r.unwrap()).collect::<Vec<TempRecord>>();

    {
        let mut map = phf_codegen::Map::new();
        for ref record in &records {
            map.entry(&record.zip_code[..], &record.to_syntax());
        }
        map.build(&mut file).unwrap();
    }

    write!(&mut file, ";\n").unwrap();
}
<|endoftext|>


<commit_before>#[cfg(linux)]
mod linux {
    use std::env;
    use std::fs;
    use std::path::PathBuf;

    const LAUNCH_SCRIPT: &str = "anup.sh";

    pub fn run() {
        let profile = env::var("PROFILE").unwrap();

        if profile != "release" {
            return;
        }

        let mut out_path = PathBuf::from("target");
        out_path.push(profile);
        out_path.push(LAUNCH_SCRIPT);

        fs::copy(LAUNCH_SCRIPT, out_path).unwrap();
    }
}

fn main() {
    #[cfg(linux)]
    linux::run();
}
<commit_msg>[build] Copy launch script on all build profiles<commit_after>#[cfg(linux)]
mod linux {
    use std::env;
    use std::fs;
    use std::path::PathBuf;

    const LAUNCH_SCRIPT: &str = "anup.sh";

    pub fn run() {
        let profile = env::var("PROFILE").unwrap();

        let mut out_path = PathBuf::from("target");
        out_path.push(profile);
        out_path.push(LAUNCH_SCRIPT);

        fs::copy(LAUNCH_SCRIPT, out_path).unwrap();
    }
}

fn main() {
    #[cfg(linux)]
    linux::run();
}
<|endoftext|>


<commit_before>use core::{mem, str};

use pretty;

#[derive(Copy, Clone)]
pub struct Buffer {
    bytes: [u8; 24],
}

impl Buffer {
    pub fn new() -> Self {
        Buffer {
            bytes: unsafe { mem::uninitialized() },
        }
    }

    pub fn format<F: Float>(&mut self, f: F) -> &str {
        f.write_to_ryu_buffer(self)
    }
}

impl Default for Buffer {
    fn default() -> Self {
        Buffer::new()
    }
}

pub trait Float: Sealed {
    #[doc(hidden)]
    fn write_to_ryu_buffer(self, buffer: &mut Buffer) -> &str;
}

impl Float for f32 {
    fn write_to_ryu_buffer(self, buffer: &mut Buffer) -> &str {
        unsafe {
            let n = pretty::f2s_buffered_n(self, &mut buffer.bytes[0]);
            str::from_utf8_unchecked(&buffer.bytes[..n])
        }
    }
}

impl Float for f64 {
    fn write_to_ryu_buffer(self, buffer: &mut Buffer) -> &str {
        unsafe {
            let n = pretty::d2s_buffered_n(self, &mut buffer.bytes[0]);
            str::from_utf8_unchecked(&buffer.bytes[..n])
        }
    }
}

pub trait Sealed {}
impl Sealed for f32 {}
impl Sealed for f64 {}
<commit_msg>Inline the buffer methods<commit_after>use core::{mem, str};

use pretty;

#[derive(Copy, Clone)]
pub struct Buffer {
    bytes: [u8; 24],
}

impl Buffer {
    #[inline]
    pub fn new() -> Self {
        Buffer {
            bytes: unsafe { mem::uninitialized() },
        }
    }

    pub fn format<F: Float>(&mut self, f: F) -> &str {
        f.write_to_ryu_buffer(self)
    }
}

impl Default for Buffer {
    #[inline]
    fn default() -> Self {
        Buffer::new()
    }
}

pub trait Float: Sealed {
    #[doc(hidden)]
    fn write_to_ryu_buffer(self, buffer: &mut Buffer) -> &str;
}

impl Float for f32 {
    #[inline]
    fn write_to_ryu_buffer(self, buffer: &mut Buffer) -> &str {
        unsafe {
            let n = pretty::f2s_buffered_n(self, &mut buffer.bytes[0]);
            str::from_utf8_unchecked(&buffer.bytes[..n])
        }
    }
}

impl Float for f64 {
    #[inline]
    fn write_to_ryu_buffer(self, buffer: &mut Buffer) -> &str {
        unsafe {
            let n = pretty::d2s_buffered_n(self, &mut buffer.bytes[0]);
            str::from_utf8_unchecked(&buffer.bytes[..n])
        }
    }
}

pub trait Sealed {}
impl Sealed for f32 {}
impl Sealed for f64 {}
<|endoftext|>


<commit_before>//! Belief Maintenance System for the agent system.
//!
//! This module adds methods and types for:
//! 1) Recording how a belief came to existence to the agent.
//! 2) Detecting inconsistences between new and old beliefs.
//! 3) Fixing those inconsitences.

use super::{ProofResult, Representation};
use super::kb::QueryInput;
use lang::{Date, Grounded};

use chrono::UTC;

use std::mem;
use std::sync::RwLock;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

#[derive(Debug)]
pub struct BmsWrapper {
    records: RwLock<Vec<*mut BmsRecord>>,
    read: AtomicBool,
    pub overwrite: RwLock<bool>,
}

impl BmsWrapper {
    pub fn new(overwrite: bool) -> BmsWrapper {
        let wrapper = BmsWrapper {
            records: RwLock::new(vec![]),
            read: AtomicBool::new(true),
            overwrite: RwLock::new(overwrite),
        };
        wrapper
    }

    pub fn new_record(&self, date: Option<Date>, value: Option<f32>) {
        let date = match date {
            Some(date) => date,
            None => UTC::now(),
        };
        let record = BmsRecord {
            locked: AtomicBool::new(false),
            produced: vec![],
            date: date,
            value: value,
            refcnt: AtomicUsize::new(1),
        };
        let raw_rec = Box::into_raw(Box::new(record));
        let mut records = self.records.write().unwrap();
        // if there is a previous record, lock it
        if records.len() > 0 {
            let raw = records.last_mut().unwrap();
            let last = unsafe { &mut **raw as &mut BmsRecord };
            last.lock();
        }
        records.push(raw_rec);
    }

    pub fn update(&self, agent: &Representation, data: &BmsWrapper) {
        let ask_processed = |entry: &(Grounded, Option<f32>), cmp_rec: &Date| {
            match *entry {
                (Grounded::Function(ref func), ..) => {
                    // check if indeed the value was produced by this producer or is more recent
                    let mut ask = false;
                    {
                        let lock = func.bms.records.read().unwrap();
                        let last = unsafe { &**lock.last().unwrap() as &BmsRecord };
                        if last.date > *cmp_rec {
                            // if it was produced, run again a test against the kb to check
                            // if it is still valid
                            ask = true;
                        }
                    }
                    if ask {
                        agent.ask_processed(QueryInput::AskRelationalFunc(func.clone()), true);
                    }
                }
                (Grounded::Terminal(ref cls), ..) => {
                    let mut ask = false;
                    {
                        let lock = cls.bms.as_ref().unwrap().records.read().unwrap();
                        let last = unsafe { &**lock.last().unwrap() as &BmsRecord };
                        if last.date > *cmp_rec {
                            ask = true;
                        }
                    }
                    if ask {
                        agent.ask_processed(QueryInput::AskClassMember(cls.clone()), true);
                    }
                }
            }
        };

        // safety lock to avoid duplicate processing of the same data
        if !data.read.load(Ordering::Acquire) {
            return;
        } else {
            data.read.store(false, Ordering::Release);
        }

        if *data.overwrite.read().unwrap() {
            let mut recs: Vec<*mut BmsRecord> = vec![];
            for rec in &*data.records.read().unwrap() {
                recs.push(rec.clone())
            }
            {
                let new_rec = &mut *self.records.write().unwrap();
                let pre_rec = &mut recs;
                mem::swap(pre_rec, new_rec);
                let last = unsafe { &mut **pre_rec.last_mut().unwrap() as &mut BmsRecord };
                last.lock();
            }
            for rec in recs {
                unsafe {
                    let r =  &*rec as &BmsRecord;
                    for entry in r.get_old_entries() {
                        ask_processed(entry, &r.date);
                    }
                    // if the record is not owned by an other object, it's safe to drop
                    if r.refcnt.load(Ordering::SeqCst) == 1 {
                        Box::from_raw(rec);
                    }
                }
            }
            return;
        }

        // check if there are any inconsistencies with the knowledge produced with
        // the previous value
        let update_rec = unsafe {
            let up_lock = data.records.read().unwrap();
            &**up_lock.last().unwrap() as &BmsRecord
        };
        // create a new record with the new data and lock the last one
        self.new_record(Some(update_rec.date.clone()), update_rec.value.clone());
        // get a reference to the last record before the new one was inserted
        let last_record = unsafe {
            let lock = &*self.records.read().unwrap();
            let l = lock.len() - 2;
            &**lock.get(l).unwrap() as &BmsRecord
        };
        if (update_rec.date > last_record.date) && (update_rec.value != last_record.value) {
            // new value is more recent, check only the last produced values and
            // append a new entry to the end of the record
            for entry in last_record.get_old_entries() {
                ask_processed(entry, &last_record.date);
            }
        } else if (update_rec.date > last_record.date) && (update_rec.value != last_record.value) {
            // new value is older, in face of new information all previously
            // produced knowledge must be checked to see if it still holds true
        } else if update_rec.value != last_record.value {
            // if both dates are the same there is an incongruency
            // replace previous record value and check that all produced knowledge
            // with that value still holds true
        }
    }

    pub fn update_producers(&self, owner: Grounded, context: &ProofResult) {
        // add the produced knowledge to each producer in case it comes
        // from a logic sentence resolution
        for a in &context.antecedents {
            match *a {
                Grounded::Terminal(ref cls) => {
                    let value = cls.get_value();
                    cls.bms
                        .as_ref()
                        .unwrap()
                        .add_entry(owner.clone(), value);
                }
                Grounded::Function(ref func) => {
                    let value = func.get_value();
                    func.bms.add_entry(owner.clone(), value);
                }
            }
        }
    }

    fn add_entry(&self, produced: Grounded, with_val: Option<f32>) {
        let mut records = self.records.write().unwrap();
        let raw = records.last_mut().unwrap();
        let record = unsafe { &mut **raw as &mut BmsRecord };
        record.add_entry((produced, with_val));
    }

    pub fn overwrite_data(&self, other: &BmsWrapper) {
        let mut lock = &mut *self.records.write().unwrap();
        // drop previous records
        for rec in lock.drain(..) {
            unsafe {
                let r = &*rec as &BmsRecord;
                if r.refcnt.fetch_sub(1, Ordering::SeqCst) == 0 {
                    Box::from_raw(rec);
                }
            }
        }
        // insert new records
        for rec in &*other.records.read().unwrap() {
            lock.push(rec.clone())
        }
        let mut lock = &mut *self.overwrite.write().unwrap();
        *lock = other.overwrite.read().unwrap().clone();
    }

    pub fn record_len(&self) -> usize {
        self.records.read().unwrap().len()
    }

    pub fn get_last_date(&self) -> &Date {
        let rec = unsafe {
            let records = self.records.read().unwrap();
            &**records.last().unwrap() as &BmsRecord
        };
        &rec.date
    }

    pub fn replace_last_val(&self, val: Option<f32>) {
        let mut records = &mut self.records.write().unwrap();
        let last = unsafe { &mut **records.last_mut().unwrap() as &mut BmsRecord };
        last.value = val;
    }
}

impl ::std::clone::Clone for BmsWrapper {
    fn clone(&self) -> BmsWrapper {
        let recs = self.records.read().unwrap();
        let mut data = vec![];
        for rec in &*recs {
            unsafe {
                let r = &**rec as &BmsRecord;
                r.refcnt.fetch_add(1, Ordering::SeqCst);
            }
            data.push(rec.clone());
        }
        let ow_value = *self.overwrite.read().unwrap();
        BmsWrapper {
            read: AtomicBool::new(self.read.load(Ordering::Acquire)),
            records: RwLock::new(data),
            overwrite: RwLock::new(ow_value),
        }
    }
}

impl ::std::ops::Drop for BmsWrapper {
    fn drop(&mut self) {
        unsafe {
            let mut lock = &mut *self.records.write().unwrap();
            // drop previous records
            for rec in lock.drain(..) {
                let r = &*rec as &BmsRecord;
                if r.refcnt.fetch_sub(1, Ordering::SeqCst) == 0 {
                    Box::from_raw(rec);
                }
            }
        }
    }
}

#[derive(Debug)]
struct BmsRecord {
    locked: AtomicBool,
    produced: Vec<(Grounded, Option<f32>)>,
    date: Date,
    value: Option<f32>,
    refcnt: AtomicUsize,
}

impl BmsRecord {
    fn add_entry(&mut self, entry: (Grounded, Option<f32>)) {
        self.produced.push(entry);
    }

    fn get_old_entries(&self) -> BmsRecIterator {
        if self.locked.load(Ordering::Relaxed) == true {
            BmsRecIterator::new(&self.produced)
        } else {
            panic!()
        }
    }

    fn lock(&mut self) {
        *self.locked.get_mut() = true;
    }
}

struct BmsRecIterator<'a> {
    cnt: usize,
    data: &'a Vec<(Grounded, Option<f32>)>,
}

impl<'a> BmsRecIterator<'a> {
    fn new(data: &'a Vec<(Grounded, Option<f32>)>) -> BmsRecIterator<'a> {
        BmsRecIterator {
            data: data,
            cnt: 0,
        }
    }
}

impl<'a> ::std::iter::Iterator for BmsRecIterator<'a> {
    type Item = &'a (Grounded, Option<f32>);

    fn next(&mut self) -> Option<&'a (Grounded, Option<f32>)> {
        let out = self.data.get(self.cnt);
        self.cnt += 1;
        out
    }
}

#[cfg(test)]
mod test {}
<commit_msg>added more bms unit tests<commit_after>//! Belief Maintenance System for the agent system.
//!
//! This module adds methods and types for:
//! 1) Recording how a belief came to existence to the agent.
//! 2) Detecting inconsistences between new and old beliefs.
//! 3) Fixing those inconsitences.

use super::{ProofResult, Representation};
use super::kb::QueryInput;
use lang::{Date, Grounded};

use chrono::UTC;

use std::mem;
use std::sync::RwLock;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

#[derive(Debug)]
pub struct BmsWrapper {
    records: RwLock<Vec<*mut BmsRecord>>,
    read: AtomicBool,
    pub overwrite: RwLock<bool>,
}

impl BmsWrapper {
    pub fn new(overwrite: bool) -> BmsWrapper {
        let wrapper = BmsWrapper {
            records: RwLock::new(vec![]),
            read: AtomicBool::new(true),
            overwrite: RwLock::new(overwrite),
        };
        wrapper
    }

    pub fn new_record(&self, date: Option<Date>, value: Option<f32>) {
        let date = match date {
            Some(date) => date,
            None => UTC::now(),
        };
        let record = BmsRecord {
            locked: AtomicBool::new(false),
            produced: vec![],
            date: date,
            value: value,
            refcnt: AtomicUsize::new(1),
        };
        let raw_rec = Box::into_raw(Box::new(record));
        let mut records = self.records.write().unwrap();
        // if there is a previous record, lock it
        if records.len() > 0 {
            let raw = records.last_mut().unwrap();
            let last = unsafe { &mut **raw as &mut BmsRecord };
            last.lock();
        }
        records.push(raw_rec);
    }

    pub fn update(&self, agent: &Representation, data: &BmsWrapper) {
        let ask_processed = |entry: &(Grounded, Option<f32>), cmp_rec: &Date| {
            match *entry {
                (Grounded::Function(ref func), ..) => {
                    // check if indeed the value was produced by this producer or is more recent
                    let mut ask = false;
                    {
                        let lock = func.bms.records.read().unwrap();
                        let last = unsafe { &**lock.last().unwrap() as &BmsRecord };
                        if last.date > *cmp_rec {
                            // if it was produced, run again a test against the kb to check
                            // if it is still valid
                            ask = true;
                        }
                    }
                    if ask {
                        agent.ask_processed(QueryInput::AskRelationalFunc(func.clone()), true);
                    }
                }
                (Grounded::Terminal(ref cls), ..) => {
                    let mut ask = false;
                    {
                        let lock = cls.bms.as_ref().unwrap().records.read().unwrap();
                        let last = unsafe { &**lock.last().unwrap() as &BmsRecord };
                        if last.date > *cmp_rec {
                            ask = true;
                        }
                    }
                    if ask {
                        agent.ask_processed(QueryInput::AskClassMember(cls.clone()), true);
                    }
                }
            }
        };

        // safety lock to avoid duplicate processing of the same data
        if !data.read.load(Ordering::Acquire) {
            return;
        } else {
            data.read.store(false, Ordering::Release);
        }

        if *data.overwrite.read().unwrap() {
            let mut recs: Vec<*mut BmsRecord> = vec![];
            for rec in &*data.records.read().unwrap() {
                recs.push(rec.clone())
            }
            {
                let new_rec = &mut *self.records.write().unwrap();
                let pre_rec = &mut recs;
                mem::swap(pre_rec, new_rec);
                let last = unsafe { &mut **pre_rec.last_mut().unwrap() as &mut BmsRecord };
                last.lock();
            }
            for rec in recs {
                unsafe {
                    let r =  &*rec as &BmsRecord;
                    for entry in r.get_old_entries() {
                        ask_processed(entry, &r.date);
                    }
                    // if the record is not owned by an other object, it's safe to drop
                    if r.refcnt.load(Ordering::SeqCst) == 1 {
                        Box::from_raw(rec);
                    }
                }
            }
            return;
        }

        // check if there are any inconsistencies with the knowledge produced with
        // the previous value
        let update_rec = unsafe {
            let up_lock = data.records.read().unwrap();
            &**up_lock.last().unwrap() as &BmsRecord
        };
        // create a new record with the new data and lock the last one
        self.new_record(Some(update_rec.date.clone()), update_rec.value.clone());
        // get a reference to the last record before the new one was inserted
        let last_record = unsafe {
            let lock = &*self.records.read().unwrap();
            let l = lock.len() - 2;
            &**lock.get(l).unwrap() as &BmsRecord
        };
        if (update_rec.date > last_record.date) && (update_rec.value != last_record.value) {
            // new value is more recent, check only the last produced values and
            // append a new entry to the end of the record
            for entry in last_record.get_old_entries() {
                ask_processed(entry, &last_record.date);
            }
        } else if (update_rec.date > last_record.date) && (update_rec.value != last_record.value) {
            // new value is older, in face of new information all previously
            // produced knowledge must be checked to see if it still holds true
        } else if update_rec.value != last_record.value {
            // if both dates are the same there is an incongruency
            // replace previous record value and check that all produced knowledge
            // with that value still holds true
        }
    }

    pub fn update_producers(&self, owner: Grounded, context: &ProofResult) {
        // add the produced knowledge to each producer in case it comes
        // from a logic sentence resolution
        for a in &context.antecedents {
            match *a {
                Grounded::Terminal(ref cls) => {
                    let value = cls.get_value();
                    cls.bms
                        .as_ref()
                        .unwrap()
                        .add_entry(owner.clone(), value);
                }
                Grounded::Function(ref func) => {
                    let value = func.get_value();
                    func.bms.add_entry(owner.clone(), value);
                }
            }
        }
    }

    fn add_entry(&self, produced: Grounded, with_val: Option<f32>) {
        let mut records = self.records.write().unwrap();
        let raw = records.last_mut().unwrap();
        let record = unsafe { &mut **raw as &mut BmsRecord };
        record.add_entry((produced, with_val));
    }

    pub fn overwrite_data(&self, other: &BmsWrapper) {
        let mut lock = &mut *self.records.write().unwrap();
        // drop previous records
        for rec in lock.drain(..) {
            unsafe {
                let r = &*rec as &BmsRecord;
                if r.refcnt.fetch_sub(1, Ordering::SeqCst) == 0 {
                    Box::from_raw(rec);
                }
            }
        }
        // insert new records
        for rec in &*other.records.read().unwrap() {
            lock.push(rec.clone())
        }
        let mut lock = &mut *self.overwrite.write().unwrap();
        *lock = other.overwrite.read().unwrap().clone();
    }

    pub fn record_len(&self) -> usize {
        self.records.read().unwrap().len()
    }

    pub fn get_last_date(&self) -> &Date {
        let rec = unsafe {
            let records = self.records.read().unwrap();
            &**records.last().unwrap() as &BmsRecord
        };
        &rec.date
    }

    pub fn replace_last_val(&self, val: Option<f32>) {
        let mut records = &mut self.records.write().unwrap();
        let last = unsafe { &mut **records.last_mut().unwrap() as &mut BmsRecord };
        last.value = val;
    }
}

impl ::std::clone::Clone for BmsWrapper {
    fn clone(&self) -> BmsWrapper {
        let recs = self.records.read().unwrap();
        let mut data = vec![];
        for rec in &*recs {
            unsafe {
                let r = &**rec as &BmsRecord;
                r.refcnt.fetch_add(1, Ordering::SeqCst);
            }
            data.push(rec.clone());
        }
        let ow_value = *self.overwrite.read().unwrap();
        BmsWrapper {
            read: AtomicBool::new(self.read.load(Ordering::Acquire)),
            records: RwLock::new(data),
            overwrite: RwLock::new(ow_value),
        }
    }
}

impl ::std::ops::Drop for BmsWrapper {
    fn drop(&mut self) {
        unsafe {
            let mut lock = &mut *self.records.write().unwrap();
            // drop previous records
            for rec in lock.drain(..) {
                let r = &*rec as &BmsRecord;
                if r.refcnt.fetch_sub(1, Ordering::SeqCst) == 0 {
                    Box::from_raw(rec);
                }
            }
        }
    }
}

#[derive(Debug)]
struct BmsRecord {
    locked: AtomicBool,
    produced: Vec<(Grounded, Option<f32>)>,
    date: Date,
    value: Option<f32>,
    refcnt: AtomicUsize,
}

impl BmsRecord {
    fn add_entry(&mut self, entry: (Grounded, Option<f32>)) {
        self.produced.push(entry);
    }

    fn get_old_entries(&self) -> BmsRecIterator {
        if self.locked.load(Ordering::Relaxed) == true {
            BmsRecIterator::new(&self.produced)
        } else {
            panic!()
        }
    }

    fn lock(&mut self) {
        *self.locked.get_mut() = true;
    }
}

struct BmsRecIterator<'a> {
    cnt: usize,
    data: &'a Vec<(Grounded, Option<f32>)>,
}

impl<'a> BmsRecIterator<'a> {
    fn new(data: &'a Vec<(Grounded, Option<f32>)>) -> BmsRecIterator<'a> {
        BmsRecIterator {
            data: data,
            cnt: 0,
        }
    }
}

impl<'a> ::std::iter::Iterator for BmsRecIterator<'a> {
    type Item = &'a (Grounded, Option<f32>);

    fn next(&mut self) -> Option<&'a (Grounded, Option<f32>)> {
        let out = self.data.get(self.cnt);
        self.cnt += 1;
        out
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn bms_rollback() {
        let rep = Representation::new();

        let fol = String::from("
            (ugly[$Pancho,u=0])
            (dog[$Pancho,u=1])
            (meat[$M1,u=1])
            (fn::eat[$M1,u=1;$Pancho])

            ((let x y)
             ((dog[x,u=1] && meat[y,u=1] && fn::eat[y,u=1;x])
              |> fat[x,u=1]))

            ((let x)
             ((fat[x,u=1] && dog[x,u=1]) |> (ugly[x,u=1] && sad[x,u=1])))
        ");
        rep.tell(fol).unwrap();
        let answ = rep.ask("(fat[$Pancho,u=1] && sad[$Pancho,u=1])".to_string());
        assert_eq!(answ.get_results_single(), Some(true));

        let fol = String::from("
            (run[$Pancho,u=1])
            ((let x)
             ((run[x,u=1] && dog[x,u=1]) |> fat[x,u=0]))
        ");
        rep.tell(fol).unwrap();
        let answ0 = rep.ask("(fat[$Pancho,u=0] && ugly[$Pancho,u=0])".to_string());
        assert_eq!(answ0.get_results_single(), Some(true));
        let answ1 = rep.ask("(sad[$Pancho,u=0])".to_string());
        assert_eq!(answ1.get_results_single(), None);
    }

    #[test]
    fn bms_review_after_change() {
        let rep = Representation::new();

        let fol = String::from("
            ( meat[$M1,u=1] )
            ( dog[$Pancho,u=1] )
            ( fn::eat[$M1,u=1;$Pancho] )
            ( ( let x, y )
              ( ( dog[x,u=1] && meat[y,u=1] && fn::eat[y,u=1;x] )
                |> fat[x,u=1] ) )
        ");
        rep.tell(fol).unwrap();
        let answ = rep.ask("(fat[$Pancho,u=1])".to_string());
        assert_eq!(answ.get_results_single(), Some(true));

        let fol = String::from("
            ( run[$Pancho,u=1] )
            (( let x ) (( dog[x,u=1] && run[x,u=1] ) |> fat[x,u=0] ))
        ");
        rep.tell(fol).unwrap();
        let answ = rep.ask("(fat[$Pancho,u=0])".to_string());
        assert_eq!(answ.get_results_single(), Some(true));

        rep.tell("(fn::eat[$M1,u=1;$Pancho])".to_string()).unwrap();
        let answ = rep.ask("(fat[$Pancho,u=1])".to_string());
        assert_eq!(answ.get_results_single(), Some(true));
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! An Fleetspeak client connector library.
//!
//! This library exposes a set of functions for writing client-side Fleetspeak
//! services. Each of these functions operates on a global connection object
//! that is lazily established. If this global connection cannot be established,
//! the library will panic (because without this connection Fleetspeak will shut
//! the service down anyway).
//!
//! Note that each service should send startup information upon its inception
//! and continue to heartbeat from time to time to notify the Fleetspeak client
//! that it did not get stuck.

mod connection;
mod error;

use std::fs::File;
use std::sync::Mutex;
use std::time::Duration;

use lazy_static::lazy_static;

pub use self::connection::Packet;
pub use self::error::{ReadError, WriteError};

/// Sends a heartbeat signal to the Fleetspeak client.
///
/// All client services should heartbeat from time to time. Otherwise, from the
/// Fleetspeak perspective, the service is unresponsive and should be restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat() -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::heartbeat(buf))
}

/// Sends a system message with startup information to the Fleetspeak client.
///
/// All clients are required to send this information on startup. If the client
/// does not receive this information quickly enough, the service will be
/// killed.
///
/// The `version` string should contain a self-reported version of the service.
/// This data is used primarily for statistics.
pub fn startup(version: &str) -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::startup(buf, version))
}

/// Sends the message to the Fleetspeak server.
///
/// The message is delivered to the server-side service as specified by the
/// packet and optionally tagged with a type if specified. This optional message
/// type is irrelevant for Fleetspeak but might be useful for the service the
/// message is delivered to.
///
/// In case of any I/O failure or malformed message (e.g. due to encoding
/// problems), an error is reported.
pub fn send<M>(packet: Packet<M>) -> Result<(), WriteError>
where
    M: prost::Message,
{
    locked(&CONNECTION.output, |buf| self::connection::send(buf, packet))
}

/// Receives the message from the Fleetspeak server.
///
/// This function will block until there is a message to be read from the input.
/// Note that in particular it means your service will be unable to heartbeat
/// properly. If you are not expecting the message to arrive quickly, you should
/// use [`collect`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`collect`]: fn.collect.html
pub fn receive<M>() -> Result<Packet<M>, ReadError>
where
    M: prost::Message + Default,
{
    locked(&CONNECTION.input, |buf| self::connection::receive(buf))
}

/// Collects the message from the Fleetspeak server.
///
/// Unlike [`receive`], `collect` will send heartbeat signals while polling for
/// the message at the specified `rate`.
///
/// This function is useful in the main loop of your service when it is not
/// supposed to do anything until a request from the server arrives. If your
/// service is actually awaiting for a specific message to come, you should
/// use [`receive`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`receive`]: fn.receive.html
pub fn collect<M>(rate: Duration) -> Result<Packet<M>, std::io::Error>
where
    M: prost::Message + Default + 'static,
{
    let (sender, receiver) = std::sync::mpsc::channel();

    std::thread::spawn(move || {
        let packet = receive();

        // We do not care whether the packet was sent, because this call can
        // fail only if the receiver closed itself. Since the receiver loops
        // indefinitely, it closes only in case there was a heartbeat error. But
        // at that point we are no longer interested in the packet.
        let _ = sender.send(packet);
    });

    loop {
        use std::sync::mpsc::TryRecvError::*;

        // The sender will block indefinitely until there is a message to pick,
        // so the disconnection branch is not possible.
        match receiver.try_recv() {
            Ok(packet) => return Ok(packet?),
            Err(Empty) => (),
            Err(Disconnected) => panic!(),
        }

        heartbeat()?;
        std::thread::sleep(rate);
    }
}

/// A connection to the Fleetspeak client.
///
/// The connection is realized through two files (specified by descriptors given
/// by the Fleetspeak client as environment variables): input and output. Each
/// of these files is guarded by a separate mutex to allow writing (e.g. for
/// sending heartbeat signals) when another thread might be busy with reading
/// messages.
struct Connection {
    input: Mutex<File>,
    output: Mutex<File>,
}

lazy_static! {
    static ref CONNECTION: Connection = {
        let mut input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let mut output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        use self::connection::handshake;
        handshake(&mut input, &mut output).expect("handshake failure");

        Connection {
            input: Mutex::new(input),
            output: Mutex::new(output),
        }
    };
}

/// Executes the given function with a file extracted from the mutex.
///
/// It might happen that the mutex becomes poisoned and this call will panic in
/// result. This should not be a problem in practice, because mutex poisoning
/// is a result of one of the threads being aborted. In case of a such scenario,
/// it is likely the service needs to be restarted anyway.
fn locked<F, T, E>(mutex: &Mutex<File>, f: F) -> Result<T, E>
where
    F: FnOnce(&mut File) -> Result<T, E>
{
    let mut file = mutex.lock().expect("poisoned connection mutex");
    f(&mut file)
}

/// Opens a file object pointed by an environment variable.
///
/// Note that this function will panic if the environment variable `var` is not
/// a valid file descriptor (in which case the library cannot be initialized and
/// the service is unlikely to work anyway).
fn open(var: &str) -> File {
    let fd = std::env::var(var)
        .expect(&format!("invalid variable `{}`", var))
        .parse()
        .expect(&format!("failed to parse file descriptor"));

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}
<commit_msg>Provide doc examples for basic functions.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! An Fleetspeak client connector library.
//!
//! This library exposes a set of functions for writing client-side Fleetspeak
//! services. Each of these functions operates on a global connection object
//! that is lazily established. If this global connection cannot be established,
//! the library will panic (because without this connection Fleetspeak will shut
//! the service down anyway).
//!
//! Note that each service should send startup information upon its inception
//! and continue to heartbeat from time to time to notify the Fleetspeak client
//! that it did not get stuck.

mod connection;
mod error;

use std::fs::File;
use std::sync::Mutex;
use std::time::Duration;

use lazy_static::lazy_static;

pub use self::connection::Packet;
pub use self::error::{ReadError, WriteError};

/// Sends a heartbeat signal to the Fleetspeak client.
///
/// All client services should heartbeat from time to time. Otherwise, from the
/// Fleetspeak perspective, the service is unresponsive and should be restarted.
///
/// The exact frequency of the required heartbeat is defined in the service
/// configuration file.
pub fn heartbeat() -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::heartbeat(buf))
}

/// Sends a system message with startup information to the Fleetspeak client.
///
/// All clients are required to send this information on startup. If the client
/// does not receive this information quickly enough, the service will be
/// killed.
///
/// The `version` string should contain a self-reported version of the service.
/// This data is used primarily for statistics.
pub fn startup(version: &str) -> Result<(), WriteError> {
    locked(&CONNECTION.output, |buf| self::connection::startup(buf, version))
}

/// Sends the message to the Fleetspeak server.
///
/// The message is delivered to the server-side service as specified by the
/// packet and optionally tagged with a type if specified. This optional message
/// type is irrelevant for Fleetspeak but might be useful for the service the
/// message is delivered to.
///
/// In case of any I/O failure or malformed message (e.g. due to encoding
/// problems), an error is reported.
///
/// # Examples
///
/// ```no_run
/// use fleetspeak::client::Packet;
///
/// fleetspeak::client::send(Packet {
///     service: String::from("example"),
///     kind: None,
///     data: String::from("Hello, World!"),
/// }).expect("failed to send the packet");
/// ```
pub fn send<M>(packet: Packet<M>) -> Result<(), WriteError>
where
    M: prost::Message,
{
    locked(&CONNECTION.output, |buf| self::connection::send(buf, packet))
}

/// Receives the message from the Fleetspeak server.
///
/// This function will block until there is a message to be read from the input.
/// Note that in particular it means your service will be unable to heartbeat
/// properly. If you are not expecting the message to arrive quickly, you should
/// use [`collect`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`collect`]: fn.collect.html
///
/// # Examples
///
/// ```no_run
/// match fleetspeak::client::receive::<String>() {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to receive the packet: {}", error),
/// }
/// ```
pub fn receive<M>() -> Result<Packet<M>, ReadError>
where
    M: prost::Message + Default,
{
    locked(&CONNECTION.input, |buf| self::connection::receive(buf))
}

/// Collects the message from the Fleetspeak server.
///
/// Unlike [`receive`], `collect` will send heartbeat signals while polling for
/// the message at the specified `rate`.
///
/// This function is useful in the main loop of your service when it is not
/// supposed to do anything until a request from the server arrives. If your
/// service is actually awaiting for a specific message to come, you should
/// use [`receive`] instead.
///
/// In case of any I/O failure or malformed message (e.g. due to parsing issues
/// or when some fields are not being present), an error is reported.
///
/// [`receive`]: fn.receive.html
///
/// # Examples
///
/// ```no_run
/// use std::time::Duration;
///
/// match fleetspeak::client::collect::<String>(Duration::from_secs(1)) {
///     Ok(packet) => println!("Hello, {}!", packet.data),
///     Err(error) => eprintln!("failed to collected the packet: {}", error),
/// }
/// ```
pub fn collect<M>(rate: Duration) -> Result<Packet<M>, std::io::Error>
where
    M: prost::Message + Default + 'static,
{
    let (sender, receiver) = std::sync::mpsc::channel();

    std::thread::spawn(move || {
        let packet = receive();

        // We do not care whether the packet was sent, because this call can
        // fail only if the receiver closed itself. Since the receiver loops
        // indefinitely, it closes only in case there was a heartbeat error. But
        // at that point we are no longer interested in the packet.
        let _ = sender.send(packet);
    });

    loop {
        use std::sync::mpsc::TryRecvError::*;

        // The sender will block indefinitely until there is a message to pick,
        // so the disconnection branch is not possible.
        match receiver.try_recv() {
            Ok(packet) => return Ok(packet?),
            Err(Empty) => (),
            Err(Disconnected) => panic!(),
        }

        heartbeat()?;
        std::thread::sleep(rate);
    }
}

/// A connection to the Fleetspeak client.
///
/// The connection is realized through two files (specified by descriptors given
/// by the Fleetspeak client as environment variables): input and output. Each
/// of these files is guarded by a separate mutex to allow writing (e.g. for
/// sending heartbeat signals) when another thread might be busy with reading
/// messages.
struct Connection {
    input: Mutex<File>,
    output: Mutex<File>,
}

lazy_static! {
    static ref CONNECTION: Connection = {
        let mut input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let mut output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        use self::connection::handshake;
        handshake(&mut input, &mut output).expect("handshake failure");

        Connection {
            input: Mutex::new(input),
            output: Mutex::new(output),
        }
    };
}

/// Executes the given function with a file extracted from the mutex.
///
/// It might happen that the mutex becomes poisoned and this call will panic in
/// result. This should not be a problem in practice, because mutex poisoning
/// is a result of one of the threads being aborted. In case of a such scenario,
/// it is likely the service needs to be restarted anyway.
fn locked<F, T, E>(mutex: &Mutex<File>, f: F) -> Result<T, E>
where
    F: FnOnce(&mut File) -> Result<T, E>
{
    let mut file = mutex.lock().expect("poisoned connection mutex");
    f(&mut file)
}

/// Opens a file object pointed by an environment variable.
///
/// Note that this function will panic if the environment variable `var` is not
/// a valid file descriptor (in which case the library cannot be initialized and
/// the service is unlikely to work anyway).
fn open(var: &str) -> File {
    let fd = std::env::var(var)
        .expect(&format!("invalid variable `{}`", var))
        .parse()
        .expect(&format!("failed to parse file descriptor"));

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}
<|endoftext|>