<commit_before>// Copyright 2020-2021 The Propverify authors
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use log::info;
use regex::Regex;
use std::ffi::OsStr;
use std::path::PathBuf;
use structopt::StructOpt;

use inkwell::builder::Builder;
use inkwell::context::Context;
use inkwell::memory_buffer::MemoryBuffer;
use inkwell::module::Linkage;
use inkwell::module::Module;
use inkwell::types::{AnyType, FunctionType};
use inkwell::values::{AnyValue, BasicValue, BasicValueEnum};
use inkwell::values::{FunctionValue, GlobalValue, PointerValue};
use inkwell::AddressSpace;

// Command line argument parsing
#[derive(StructOpt)]
#[structopt(
    name = "rvt-patch-llvm",
    about = "Preprocess rustc generated llvm code for verification.",
    // version number is taken automatically from Cargo.toml
)]
struct Opt {
    /// Input file.
    #[structopt(name = "INPUT", parse(from_os_str))]
    input: PathBuf,

    /// Output file
    #[structopt(
        short,
        long,
        name = "OUTPUT",
        parse(from_os_str),
        default_value = "out"
    )]
    output: PathBuf,

    /// Eliminate feature tests
    #[structopt(short, long)]
    features: bool,

    /// Call initializers from main
    #[structopt(short, long)]
    initializers: bool,

    /// SeaHorn preparation (conflicts with --initializers)
    #[structopt(short, long, conflicts_with = "initializers")]
    seahorn: bool,

    /// Increase message verbosity
    #[structopt(short, long, parse(from_occurrences))]
    verbosity: usize,
}

fn main() {
    let opt = Opt::from_args();

    #[rustfmt::skip]
    stderrlog::new()
        .verbosity(opt.verbosity)
        .init()
        .unwrap();

    // Read the input file
    info!("Reading input from {}", opt.input.to_str().unwrap());
    let memory_buffer =
        MemoryBuffer::create_from_file(&opt.input).expect("ERROR: failed to open file.");

    let context = Context::create();
    let mut module = context
        .create_module_from_ir(memory_buffer)
        .expect("ERROR: failed to create module.");

    if opt.features {
        info!("Patching feature functions");
        stub_functions(
            &context,
            &module,
            &[
                "std::std_detect::detect::arch::__is_feature_detected::avx2",
                "std::std_detect::detect::arch::__is_feature_detected::ssse3",
                "std::std_detect::detect::check_for",
            ],
            |builder| {
                builder.build_return(Some(&context.bool_type().const_int(0, false)));
                ()
            },
        );
    }

    if opt.initializers {
        handle_initializers(&context, &mut module);
    }

    if opt.seahorn {
        handle_main(&module);

        handle_panic(&context, &module);

        replace_def_with_dec(
            &module,
            &Regex::new(r"^_ZN3std2io5stdio7_eprint17h[a-f0-9]{16}E$").unwrap(),
        );
        replace_def_with_dec(
            &module,
            &Regex::new(r"^_ZN3std2io5stdio6_print17h[a-f0-9]{16}E$").unwrap(),
        );
    }

    // Write output file
    info!("Writing output to {}", opt.output.to_str().unwrap());
    if opt.output.extension() == Some(OsStr::new("bc")) {
        // output bitcode
        // TODO: this function returns bool but the doc doesn't say anything about it.
        module.write_bitcode_to_path(&opt.output);
    } else {
        // output disassembled bitcode
        module
            .print_to_file(&opt.output)
            .expect("ERROR: failed to write to file.");
    }
}

////////////////////////////////////////////////////////////////
// Transformations associated with initializers
////////////////////////////////////////////////////////////////

fn handle_initializers(context: &Context, module: &mut Module) {
    if let Some(initializer) =
        collect_initializers(context, module, ".init_array", "__init_function")
    {
        info!(
            "Combined .init_array* initializers into '{}'",
            initializer.get_name().to_str().unwrap()
        );

        let main = module
            .get_function("main")
            .expect("Unable to find 'main' function");
        let mut args = get_fn_args(main);
        assert!(args.len() == 2); // We expect "i32 @main(i32 %0, i8** %1)"
        let i8_type = context.i8_type();
        let pi8_type = i8_type.ptr_type(AddressSpace::Generic);
        let ppi8_type = pi8_type.ptr_type(AddressSpace::Generic);
        args.push(ppi8_type.const_null().as_basic_value_enum());
        insert_call_at_head(context, initializer, args, main);
        info!(
            "Inserted call to '{}' into 'main'",
            initializer.get_name().to_str().unwrap()
        )
    } else {
        info!("No initializers to handle")
    }
}

/// Collect all the initializers in a section (whose name starts with 'prefix')
/// into a single function that calls all the initializers.
fn collect_initializers<'a>(
    context: &Context,
    module: &mut Module<'a>,
    prefix: &str,
    nm: &str,
) -> Option<FunctionValue<'a>> {
    let vs = collect_variables_in_section(module, prefix);
    for v in &vs {
        info!("Found initializer {:?}", v.get_name().to_str().unwrap());
    }

    let fps: Vec<PointerValue> = vs.iter().map(get_initializer_function).collect();

    if !fps.is_empty() {
        let fp = fps[0];

        // dereference the pointer type
        let fp_type = fp.get_type().get_element_type().into_function_type();
        info!("Initializer type {:?}", fp_type.print_to_string());

        Some(build_fanout(context, module, nm, fp_type, fps))
    } else {
        None
    }
}

/// Collect variables that are assigned to a section whose name matches prefix
fn collect_variables_in_section<'a>(module: &Module<'a>, prefix: &str) -> Vec<GlobalValue<'a>> {
    let mut vs = Vec::new();
    // todo: should implement an iterator for global values so that the following loop
    // becomes just iter().filter_map().filter_map().filter()
    let mut og = module.get_first_global();
    while let Some(g) = og {
        if let Some(s) = g.get_section() {
            if let Ok(s) = s.to_str() {
                if s.starts_with(prefix) {
                    vs.push(g)
                }
            }
        }
        og = g.get_next_global();
    }
    vs
}

/// Convert the contents of an initializer section to a function pointer.
///
/// Initializer sections contain structs where the first field is a function pointer cast to
/// some other type.
fn get_initializer_function<'a>(v: &GlobalValue<'a>) -> PointerValue<'a> {
    let i = v.get_initializer().unwrap().into_struct_value();
    assert!(i.get_num_operands() == 2); // expecting two fields in struct
    let i = i.get_operand(0).unwrap();
    assert!(i.get_num_operands() == 1); // expecting bitcast
    let fp = i.get_operand(0).unwrap();
    fp.into_pointer_value()
}

/// Given a list of functions of type 'ty', build a function that calls each function in order
///
/// Assumes (without checking) that return type is void
///
///     define void @fanout(i32 %0, i32 %1) {
///     entry:
///       call void @f1(i32 %0, i32 %1)
///       call void @f2(i32 %0, i32 %1)
///       call void @f3(i32 %0, i32 %1)
///       ret void
///     }
///
fn build_fanout<'a>(
    context: &Context,
    module: &mut Module<'a>,
    nm: &str,
    ty: FunctionType<'a>,
    fps: Vec<PointerValue<'a>>,
) -> FunctionValue<'a> {
    let function = module.add_function(nm, ty, None);
    let args = get_fn_args(function);
    let basic_block = context.append_basic_block(function, "entry");
    let builder = context.create_builder();
    builder.position_at_end(basic_block);

    for fp in fps {
        builder.build_call(fp, &args, "");
        builder.build_return(None);
        // println!("Built function {:?}", function)
    }

    function
}

fn get_fn_args<'a>(function: FunctionValue<'a>) -> Vec<BasicValueEnum<'a>> {
    (0..function.count_params())
        .map(|i| function.get_nth_param(i).unwrap())
        .collect()
}

fn insert_call_at_head<'a>(
    context: &Context,
    f: FunctionValue<'a>,
    args: Vec<BasicValueEnum<'a>>,
    insertee: FunctionValue<'a>,
) {
    let bb = insertee
        .get_first_basic_block()
        .expect("Unable to find function to insert function call into");
    let first_instruction = bb
        .get_first_instruction()
        .expect("Unable to find where to insert function call into function");
    let builder = context.create_builder();
    builder.position_before(&first_instruction);
    builder.build_call(f, &args, "");
}

////////////////////////////////////////////////////////////////
// Function transformation functions
////////////////////////////////////////////////////////////////

/// Apply Rust mangling rule to a function name like 'foo::bar'.
/// Does not insert a unique hash at the end or the final 'E' character.
fn rust_mangle(rust_name: &str) -> String {
    let mangled =
        &rust_name
        .split("::")
        .into_iter()
        .map(|x| format!("{}{}", x.len(), x))
        .collect::<Vec<String>>()
        .concat();
    "_ZN".to_string() + mangled
}

/// Delete the body of a function
fn delete_body<'ctx>(fun: &FunctionValue<'ctx>) {
    for bb in fun.get_basic_blocks() {
        unsafe {
            // Any use of `bb` after calling `delete` is unsafe
            bb.delete().expect("delete basic block");
        }
    }
}

/// Stub out a list of functions by deleting the existing
/// body of each function and using mk_stub to construct a new body.
///
/// Functions are identified specified by names like "foo::bar"
/// and exclude the uniquifying hash code at the end.
fn stub_functions<F>(context: &Context, module: &Module, functions: &[&str], mk_stub: F)
where
    F: Fn(&Builder),
{
    let builder = context.create_builder();
    for function in functions {
        for fv in get_function_by_unmangled_name(&module, function) {
            info!("Stubbing out function function {}", function);
            delete_body(&fv);
            let basic_block = context.append_basic_block(fv, "entry");
            builder.position_at_end(basic_block);
            mk_stub(&builder)
        }
    }
}

////////////////////////////////////////////////////////////////
// Transformations associated with SeaHorn
////////////////////////////////////////////////////////////////

fn handle_main(module: &Module) {
    // Remove the main function rustc generates. I don't know why, but passing
    // --entry=.. to seahorn is not enough, we also have to remove main.
    if let Some(main) = module.get_function("main") {
        unsafe {
            // Any use of `main` after calling `delete` is unsafe
            main.delete();
        }
        info!("Deleted 'main' (was added by rustc).");
    }
}

/// Find a function whose name matches regex `re`
fn get_function_by_regex<'ctx>(module: &'ctx Module, re: &Regex) -> Vec<FunctionValue<'ctx>> {
    get_function(module, |name| re.is_match(name))
}

/// Find a function whose name matches a Rust-mangled name
fn get_function_by_unmangled_name<'ctx>(module: &'ctx Module, prefix: &str) -> Vec<FunctionValue<'ctx>> {
    let prefix = format!("{}17h", rust_mangle(prefix));
    get_function(module, |name| name.starts_with(&prefix))
}

/// Find a function whose name satisfies predicate `is_match`
fn get_function<'ctx, P>(module: &'ctx Module, is_match: P) -> Vec<FunctionValue<'ctx>>
where
    P: Fn(&str) -> bool,
{
    let mut funs = vec![];
    let mut op_fun = module.get_first_function();
    while let Some(fun) = op_fun {
        if fun.get_name()
            .to_str()
            .map(&is_match)
            .unwrap_or(false)
        {
            funs.push(fun);
        }
        op_fun = fun.get_next_function();
    }
    funs
}

fn handle_panic(context: &Context, module: &Module) {
    // TODO: make "spanic" a CL arg.
    if let Some(spanic) = module.get_function("spanic") {
        let builder = context.create_builder();

        // Delete the body of panic functions, and replace it with a call to
        // `spanic`.

        // Note that std::panicking::begin_panic can have multiple instances
        // with different hash suffix, I'm not sure why.
        for fv in module.get_function("rust_begin_unwind").into_iter()
            .chain(get_function_by_unmangled_name(&module, "std::panicking::begin_panic"))
            .chain(get_function_by_unmangled_name(&module, "core::panicking::panic"))
        {
            delete_body(&fv);
            let basic_block = context.append_basic_block(fv, "entry");
            builder.position_at_end(basic_block);
            builder.build_call(spanic, &[], "call");
            builder.build_return(None);

            info!("Replaced the body of '{}' with a call to '{}'.",
                  fv.get_name().to_string_lossy(),
                  spanic.get_name().to_string_lossy(),
            );
        }
    }
}

/// Change a function to a declaration by
/// deleting all basic blocks and modifying metadata
/// such as personality_function, linkage, etc.
fn replace_def_with_dec(module: &Module, re: &Regex) {
    for fun in get_function_by_regex(module, re) {
        delete_body(&fun);
        fun.remove_personality_function();
        fun.set_linkage(Linkage::External);
        info!(
            "Removed the implementation of '{}'.",
            fun.get_name().to_str().unwrap()
        );
    }
}

////////////////////////////////////////////////////////////////
// End
////////////////////////////////////////////////////////////////
<commit_msg>Patch-LLVM: Add --intrinsics<commit_after>// Copyright 2020-2021 The Propverify authors
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use log::info;
use regex::Regex;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::path::PathBuf;
use structopt::StructOpt;

use inkwell::builder::Builder;
use inkwell::context::Context;
use inkwell::memory_buffer::MemoryBuffer;
use inkwell::module::Linkage;
use inkwell::module::Module;
use inkwell::types::{AnyType, FunctionType};
use inkwell::values::{AnyValue, BasicValue, BasicValueEnum};
use inkwell::values::{FunctionValue, GlobalValue, PointerValue};
use inkwell::AddressSpace;

// Command line argument parsing
#[derive(StructOpt)]
#[structopt(
    name = "rvt-patch-llvm",
    about = "Preprocess rustc generated llvm code for verification.",
    // version number is taken automatically from Cargo.toml
)]
struct Opt {
    /// Input file.
    #[structopt(name = "INPUT", parse(from_os_str))]
    input: PathBuf,

    /// Output file
    #[structopt(
        short,
        long,
        name = "OUTPUT",
        parse(from_os_str),
        default_value = "out"
    )]
    output: PathBuf,

    /// Eliminate feature tests
    #[structopt(long)]
    features: bool,

    /// Call initializers from main
    #[structopt(long)]
    initializers: bool,

    /// Substitute function calls for SIMD intrinsic operations
    #[structopt(long)]
    intrinsics: bool,

    /// SeaHorn preparation (conflicts with --initializers)
    #[structopt(long, conflicts_with = "initializers")]
    seahorn: bool,

    /// Increase message verbosity
    #[structopt(short, long, parse(from_occurrences))]
    verbosity: usize,
}

fn main() {
    let opt = Opt::from_args();

    #[rustfmt::skip]
    stderrlog::new()
        .verbosity(opt.verbosity)
        .init()
        .unwrap();

    // Read the input file
    info!("Reading input from {}", opt.input.to_str().unwrap());
    let memory_buffer =
        MemoryBuffer::create_from_file(&opt.input).expect("ERROR: failed to open file.");

    let context = Context::create();
    let mut module = context
        .create_module_from_ir(memory_buffer)
        .expect("ERROR: failed to create module.");

    if opt.features {
        info!("Patching feature functions");
        stub_functions(
            &context,
            &module,
            &[
                "std::std_detect::detect::arch::__is_feature_detected::avx2",
                "std::std_detect::detect::arch::__is_feature_detected::ssse3",
                "std::std_detect::detect::check_for",
            ],
            |builder| {
                builder.build_return(Some(&context.bool_type().const_int(0, false)));
                ()
            },
        );
    }

    if opt.initializers {
        handle_initializers(&context, &mut module);
    }

    if opt.intrinsics {
        // intrinsics
        let is = get_function(&module, |name| name.starts_with("llvm.x86"));

        // simd_emulation functions
        let rs = get_function(&module, |name| name.starts_with("llvm_x86"));
        let rs: HashMap<&str, FunctionValue> = rs.iter().filter_map(|f| f.get_name().to_str().ok().map(|nm| (nm, *f))).collect();

        info!("Found simd_emulation functions {:?}", rs.keys());

        for i in is {
            let i_name = i.get_name();
            let r_name = i_name.to_str().expect("valid UTF8 symbol name");
            let r_name: String= r_name.replace(".", "_");
            if let Some(r) = rs.get(r_name.as_str()) {
                info!("Replacing intrinsic {:?} with {}", i_name, r_name);
                i.replace_all_uses_with(*r);
            } else {
                info!("Did not find replacement for {:?}", i_name);
            }
        }
    }

    if opt.seahorn {
        handle_main(&module);

        handle_panic(&context, &module);

        replace_def_with_dec(
            &module,
            &Regex::new(r"^_ZN3std2io5stdio7_eprint17h[a-f0-9]{16}E$").unwrap(),
        );
        replace_def_with_dec(
            &module,
            &Regex::new(r"^_ZN3std2io5stdio6_print17h[a-f0-9]{16}E$").unwrap(),
        );
    }

    // Write output file
    info!("Writing output to {}", opt.output.to_str().unwrap());
    if opt.output.extension() == Some(OsStr::new("bc")) {
        // output bitcode
        // TODO: this function returns bool but the doc doesn't say anything about it.
        module.write_bitcode_to_path(&opt.output);
    } else {
        // output disassembled bitcode
        module
            .print_to_file(&opt.output)
            .expect("ERROR: failed to write to file.");
    }
}

////////////////////////////////////////////////////////////////
// Transformations associated with initializers
////////////////////////////////////////////////////////////////

fn handle_initializers(context: &Context, module: &mut Module) {
    if let Some(initializer) =
        collect_initializers(context, module, ".init_array", "__init_function")
    {
        info!(
            "Combined .init_array* initializers into '{}'",
            initializer.get_name().to_str().unwrap()
        );

        let main = module
            .get_function("main")
            .expect("Unable to find 'main' function");
        let mut args = get_fn_args(main);
        assert!(args.len() == 2); // We expect "i32 @main(i32 %0, i8** %1)"
        let i8_type = context.i8_type();
        let pi8_type = i8_type.ptr_type(AddressSpace::Generic);
        let ppi8_type = pi8_type.ptr_type(AddressSpace::Generic);
        args.push(ppi8_type.const_null().as_basic_value_enum());
        insert_call_at_head(context, initializer, args, main);
        info!(
            "Inserted call to '{}' into 'main'",
            initializer.get_name().to_str().unwrap()
        )
    } else {
        info!("No initializers to handle")
    }
}

/// Collect all the initializers in a section (whose name starts with 'prefix')
/// into a single function that calls all the initializers.
fn collect_initializers<'a>(
    context: &Context,
    module: &mut Module<'a>,
    prefix: &str,
    nm: &str,
) -> Option<FunctionValue<'a>> {
    let vs = collect_variables_in_section(module, prefix);
    for v in &vs {
        info!("Found initializer {:?}", v.get_name().to_str().unwrap());
    }

    let fps: Vec<PointerValue> = vs.iter().map(get_initializer_function).collect();

    if !fps.is_empty() {
        let fp = fps[0];

        // dereference the pointer type
        let fp_type = fp.get_type().get_element_type().into_function_type();
        info!("Initializer type {:?}", fp_type.print_to_string());

        Some(build_fanout(context, module, nm, fp_type, fps))
    } else {
        None
    }
}

/// Collect variables that are assigned to a section whose name matches prefix
fn collect_variables_in_section<'a>(module: &Module<'a>, prefix: &str) -> Vec<GlobalValue<'a>> {
    let mut vs = Vec::new();
    // todo: should implement an iterator for global values so that the following loop
    // becomes just iter().filter_map().filter_map().filter()
    let mut og = module.get_first_global();
    while let Some(g) = og {
        if let Some(s) = g.get_section() {
            if let Ok(s) = s.to_str() {
                if s.starts_with(prefix) {
                    vs.push(g)
                }
            }
        }
        og = g.get_next_global();
    }
    vs
}

/// Convert the contents of an initializer section to a function pointer.
///
/// Initializer sections contain structs where the first field is a function pointer cast to
/// some other type.
fn get_initializer_function<'a>(v: &GlobalValue<'a>) -> PointerValue<'a> {
    let i = v.get_initializer().unwrap().into_struct_value();
    assert!(i.get_num_operands() == 2); // expecting two fields in struct
    let i = i.get_operand(0).unwrap();
    assert!(i.get_num_operands() == 1); // expecting bitcast
    let fp = i.get_operand(0).unwrap();
    fp.into_pointer_value()
}

/// Given a list of functions of type 'ty', build a function that calls each function in order
///
/// Assumes (without checking) that return type is void
///
///     define void @fanout(i32 %0, i32 %1) {
///     entry:
///       call void @f1(i32 %0, i32 %1)
///       call void @f2(i32 %0, i32 %1)
///       call void @f3(i32 %0, i32 %1)
///       ret void
///     }
///
fn build_fanout<'a>(
    context: &Context,
    module: &mut Module<'a>,
    nm: &str,
    ty: FunctionType<'a>,
    fps: Vec<PointerValue<'a>>,
) -> FunctionValue<'a> {
    let function = module.add_function(nm, ty, None);
    let args = get_fn_args(function);
    let basic_block = context.append_basic_block(function, "entry");
    let builder = context.create_builder();
    builder.position_at_end(basic_block);

    for fp in fps {
        builder.build_call(fp, &args, "");
        builder.build_return(None);
        // println!("Built function {:?}", function)
    }

    function
}

fn get_fn_args<'a>(function: FunctionValue<'a>) -> Vec<BasicValueEnum<'a>> {
    (0..function.count_params())
        .map(|i| function.get_nth_param(i).unwrap())
        .collect()
}

fn insert_call_at_head<'a>(
    context: &Context,
    f: FunctionValue<'a>,
    args: Vec<BasicValueEnum<'a>>,
    insertee: FunctionValue<'a>,
) {
    let bb = insertee
        .get_first_basic_block()
        .expect("Unable to find function to insert function call into");
    let first_instruction = bb
        .get_first_instruction()
        .expect("Unable to find where to insert function call into function");
    let builder = context.create_builder();
    builder.position_before(&first_instruction);
    builder.build_call(f, &args, "");
}

////////////////////////////////////////////////////////////////
// Function transformation functions
////////////////////////////////////////////////////////////////

/// Apply Rust mangling rule to a function name like 'foo::bar'.
/// Does not insert a unique hash at the end or the final 'E' character.
fn rust_mangle(rust_name: &str) -> String {
    let mangled =
        &rust_name
        .split("::")
        .into_iter()
        .map(|x| format!("{}{}", x.len(), x))
        .collect::<Vec<String>>()
        .concat();
    "_ZN".to_string() + mangled
}

/// Delete the body of a function
fn delete_body<'ctx>(fun: &FunctionValue<'ctx>) {
    for bb in fun.get_basic_blocks() {
        unsafe {
            // Any use of `bb` after calling `delete` is unsafe
            bb.delete().expect("delete basic block");
        }
    }
}

/// Stub out a list of functions by deleting the existing
/// body of each function and using mk_stub to construct a new body.
///
/// Functions are identified specified by names like "foo::bar"
/// and exclude the uniquifying hash code at the end.
fn stub_functions<F>(context: &Context, module: &Module, functions: &[&str], mk_stub: F)
where
    F: Fn(&Builder),
{
    let builder = context.create_builder();
    for function in functions {
        for fv in get_function_by_unmangled_name(&module, function) {
            info!("Stubbing out function function {}", function);
            delete_body(&fv);
            let basic_block = context.append_basic_block(fv, "entry");
            builder.position_at_end(basic_block);
            mk_stub(&builder)
        }
    }
}

////////////////////////////////////////////////////////////////
// Transformations associated with SeaHorn
////////////////////////////////////////////////////////////////

fn handle_main(module: &Module) {
    // Remove the main function rustc generates. I don't know why, but passing
    // --entry=.. to seahorn is not enough, we also have to remove main.
    if let Some(main) = module.get_function("main") {
        unsafe {
            // Any use of `main` after calling `delete` is unsafe
            main.delete();
        }
        info!("Deleted 'main' (was added by rustc).");
    }
}

/// Find a function whose name matches regex `re`
fn get_function_by_regex<'ctx>(module: &'ctx Module, re: &Regex) -> Vec<FunctionValue<'ctx>> {
    get_function(module, |name| re.is_match(name))
}

/// Find a function whose name matches a Rust-mangled name
fn get_function_by_unmangled_name<'ctx>(module: &'ctx Module, prefix: &str) -> Vec<FunctionValue<'ctx>> {
    let prefix = format!("{}17h", rust_mangle(prefix));
    get_function(module, |name| name.starts_with(&prefix))
}

/// Find a function whose name satisfies predicate `is_match`
fn get_function<'ctx, P>(module: &'ctx Module, is_match: P) -> Vec<FunctionValue<'ctx>>
where
    P: Fn(&str) -> bool,
{
    let mut funs = vec![];
    let mut op_fun = module.get_first_function();
    while let Some(fun) = op_fun {
        if fun.get_name()
            .to_str()
            .map(&is_match)
            .unwrap_or(false)
        {
            funs.push(fun);
        }
        op_fun = fun.get_next_function();
    }
    funs
}

fn handle_panic(context: &Context, module: &Module) {
    // TODO: make "spanic" a CL arg.
    if let Some(spanic) = module.get_function("spanic") {
        let builder = context.create_builder();

        // Delete the body of panic functions, and replace it with a call to
        // `spanic`.

        // Note that std::panicking::begin_panic can have multiple instances
        // with different hash suffix, I'm not sure why.
        for fv in module.get_function("rust_begin_unwind").into_iter()
            .chain(get_function_by_unmangled_name(&module, "std::panicking::begin_panic"))
            .chain(get_function_by_unmangled_name(&module, "core::panicking::panic"))
        {
            delete_body(&fv);
            let basic_block = context.append_basic_block(fv, "entry");
            builder.position_at_end(basic_block);
            builder.build_call(spanic, &[], "call");
            builder.build_return(None);

            info!("Replaced the body of '{}' with a call to '{}'.",
                  fv.get_name().to_string_lossy(),
                  spanic.get_name().to_string_lossy(),
            );
        }
    }
}

/// Change a function to a declaration by
/// deleting all basic blocks and modifying metadata
/// such as personality_function, linkage, etc.
fn replace_def_with_dec(module: &Module, re: &Regex) {
    for fun in get_function_by_regex(module, re) {
        delete_body(&fun);
        fun.remove_personality_function();
        fun.set_linkage(Linkage::External);
        info!(
            "Removed the implementation of '{}'.",
            fun.get_name().to_str().unwrap()
        );
    }
}

////////////////////////////////////////////////////////////////
// End
////////////////////////////////////////////////////////////////
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::io::prelude::*;
// use std::io::Read;   // will include this

fn main() {
    let args: Vec<String> = env::args().collect();

    let config = Config::new(&args);

    println!("Searching for {}", config.query);
    println!("In file {}", config.filename);

    let mut f = File::open(config.filename).expect("file not found");

    let mut contents = String::new();
    f.read_to_string(&mut contents)
        .expect("somthing went wrong reading the file");

    println!("With text:\n{}", contents);
}

struct Config {
    query: String,
    filename: String,
}

impl Config {
    fn new(args: &[String]) -> Config {
        let query = args[1].clone();
        let filename = args[2].clone();

        Config {
            query: query,
            filename: filename,
        }
    }
}
<commit_msg>On io refactor: use option instead of panic.<commit_after>use std::env;
use std::process;
use std::fs::File;
use std::io::prelude::*;
// use std::io::Read;   // will include this

fn main() {
    let args: Vec<String> = env::args().collect();

    let config = Config::new(&args).unwrap_or_else(|err| {
        println!("Problem parsing arguments: {}", err);
        process::exit(1);
    });

    println!("Searching for {}", config.query);
    println!("In file {}", config.filename);

    let mut f = File::open(config.filename).expect("file not found");

    let mut contents = String::new();
    f.read_to_string(&mut contents)
        .expect("somthing went wrong reading the file");

    println!("With text:\n{}", contents);
}

struct Config {
    query: String,
    filename: String,
}

impl Config {
    fn new(args: &[String]) -> Result<Config, &'static str> {
        if args.len() < 3 {
            return Err("not enough arguments");
        }
        let query = args[1].clone();
        let filename = args[2].clone();

        Ok(Config {
            query: query,
            filename: filename,
        })
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use bazel_protos;
use boxfuture::{Boxable, BoxFuture};
use futures::Future;
use futures::future::join_all;
use hashing::{Digest, Fingerprint};
use itertools::Itertools;
use {File, PathStat, Store};
use protobuf;
use std::ffi::OsString;
use std::fmt;
use std::sync::Arc;

#[derive(Clone, PartialEq)]
pub struct Snapshot {
  // TODO: In a follow-up commit, fingerprint will be removed, and digest will be made non-optional.
  // They both exist right now as a compatibility shim so that the tar-based code and
  // Directory-based code can peacefully co-exist.
  pub fingerprint: Fingerprint,
  pub digest: Option<Digest>,
  pub path_stats: Vec<PathStat>,
}

// StoreFileByDigest allows a File to be saved to an underlying Store, in such a way that it can be
// looked up by the Digest produced by the store_by_digest method.
// It is a separate trait so that caching implementations can be written which wrap the Store (used
// to store the bytes) and VFS (used to read the files off disk if needed).
pub trait StoreFileByDigest<Error> {
  fn store_by_digest(&self, file: &File) -> BoxFuture<Digest, Error>;
}

impl Snapshot {
  pub fn from_path_stats<
    S: StoreFileByDigest<Error> + Sized + Clone,
    Error: fmt::Debug + 'static + Send,
  >(
    store: Arc<Store>,
    file_digester: S,
    mut path_stats: Vec<PathStat>,
  ) -> BoxFuture<Snapshot, String> {
    let mut file_futures: Vec<BoxFuture<bazel_protos::remote_execution::FileNode, String>> =
      Vec::new();
    let mut dir_futures: Vec<BoxFuture<bazel_protos::remote_execution::DirectoryNode, String>> =
      Vec::new();

    path_stats.sort_by(|a, b| a.path().cmp(b.path()));

    for (first_component, group) in
      &path_stats.iter().cloned().group_by(|s| {
        s.path().components().next().unwrap().as_os_str().to_owned()
      })
    {
      let mut path_group: Vec<PathStat> = group.collect();
      if path_group.len() == 1 && path_group.get(0).unwrap().path().components().count() == 1 {
        // Exactly one entry with exactly one component indicates either a file in this directory,
        // or an empty directory.
        // If the child is a non-empty directory, or a file therein, there must be multiple
        // PathStats with that prefix component, and we will handle that in the recursive
        // save_directory call.

        match path_group.pop().unwrap() {
          PathStat::File { ref stat, .. } => {
            let is_executable = stat.is_executable;
            file_futures.push(
              file_digester
                .clone()
                .store_by_digest(&stat)
                .map_err(|e| format!("{:?}", e))
                .and_then(move |digest| {
                  let mut file_node = bazel_protos::remote_execution::FileNode::new();
                  file_node.set_name(osstring_as_utf8(first_component)?);
                  file_node.set_digest((&digest).into());
                  file_node.set_is_executable(is_executable);
                  Ok(file_node)
                })
                .to_boxed(),
            );
          }
          PathStat::Dir { .. } => {
            // Because there are no children of this Dir, it must be empty.
            dir_futures.push(
              store
                .record_directory(&bazel_protos::remote_execution::Directory::new(), true)
                .map(move |digest| {
                  let mut directory_node = bazel_protos::remote_execution::DirectoryNode::new();
                  directory_node.set_name(osstring_as_utf8(first_component).unwrap());
                  directory_node.set_digest((&digest).into());
                  directory_node
                })
                .to_boxed(),
            );
          }
        }
      } else {
        dir_futures.push(
          // TODO: Memoize this in the graph
          Snapshot::from_path_stats(
            store.clone(),
            file_digester.clone(),
            paths_of_child_dir(path_group),
          ).and_then(move |snapshot| {
            let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
            dir_node.set_name(osstring_as_utf8(first_component)?);
            dir_node.set_digest((&snapshot.digest.unwrap()).into());
            Ok(dir_node)
          })
            .to_boxed(),
        );
      }
    }
    join_all(dir_futures)
      .join(join_all(file_futures))
      .and_then(move |(dirs, files)| {
        let mut directory = bazel_protos::remote_execution::Directory::new();
        directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
        directory.set_files(protobuf::RepeatedField::from_vec(files));
        store.record_directory(&directory, true).map(move |digest| {
          Snapshot {
            fingerprint: digest.0,
            digest: Some(digest),
            path_stats: path_stats,
          }
        })
      })
      .to_boxed()
  }
}

impl fmt::Debug for Snapshot {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(
      f,
      "Snapshot({}, digest={:?}, entries={})",
      self.fingerprint.to_hex(),
      self.digest,
      self.path_stats.len()
    )
  }
}

fn paths_of_child_dir(paths: Vec<PathStat>) -> Vec<PathStat> {
  paths
    .into_iter()
    .filter_map(|s| {
      if s.path().components().count() == 1 {
        return None;
      }
      Some(match s {
        PathStat::File { path, stat } => {
          PathStat::File {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
        PathStat::Dir { path, stat } => {
          PathStat::Dir {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
      })
    })
    .collect()
}

fn osstring_as_utf8(path: OsString) -> Result<String, String> {
  path.into_string().map_err(|p| {
    format!("{:?}'s file_name is not representable in UTF8", p)
  })
}

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

  use boxfuture::{BoxFuture, Boxable};
  use futures::future::Future;
  use hashing::{Digest, Fingerprint};
  use tempdir::TempDir;
  use self::testutil::make_file;

  use super::super::{File, PathGlobs, PathStat, PosixFS, ResettablePool, Snapshot, Store,
                     StoreFileByDigest, VFS};

  use std;
  use std::error::Error;
  use std::path::PathBuf;
  use std::sync::Arc;

  const STR: &str = "European Burmese";

  fn setup() -> (Arc<Store>, TempDir, Arc<PosixFS>, FileSaver) {
    let pool = Arc::new(ResettablePool::new("test-pool-".to_string()));
    // TODO: Pass a remote CAS address through.
    let store = Arc::new(
      Store::local_only(TempDir::new("lmdb_store").unwrap(), pool.clone()).unwrap(),
    );
    let dir = TempDir::new("root").unwrap();
    let posix_fs = Arc::new(PosixFS::new(dir.path(), pool, vec![]).unwrap());
    let file_saver = FileSaver(store.clone(), posix_fs.clone());
    (store, dir, posix_fs, file_saver)
  }

  #[test]
  fn snapshot_one_file() {
    let (store, dir, posix_fs, digester) = setup();

    let file_name = PathBuf::from("roland");
    make_file(&dir.path().join(&file_name), STR.as_bytes(), 0o600);

    let path_stats = expand_all_sorted(posix_fs);
    // TODO: Inline when only used once
    let fingerprint = Fingerprint::from_hex_string(
      "63949aa823baf765eff07b946050d76ec0033144c785a94d3ebd82baa931cd16",
    ).unwrap();
    assert_eq!(
      Snapshot::from_path_stats(store, digester, path_stats.clone())
        .wait()
        .unwrap(),
      Snapshot {
        fingerprint: fingerprint,
        digest: Some(Digest(fingerprint, 80)),
        path_stats: path_stats,
      }
    );
  }

  #[test]
  fn snapshot_recursive_directories() {
    let (store, dir, posix_fs, digester) = setup();

    let cats = PathBuf::from("cats");
    let roland = cats.join("roland");
    std::fs::create_dir_all(&dir.path().join(cats)).unwrap();
    make_file(&dir.path().join(&roland), STR.as_bytes(), 0o600);

    let path_stats = expand_all_sorted(posix_fs);
    // TODO: Inline when only used once
    let fingerprint = Fingerprint::from_hex_string(
      "8b1a7ea04eaa2527b35683edac088bc826117b53b7ec6601740b55e20bce3deb",
    ).unwrap();
    assert_eq!(
      Snapshot::from_path_stats(store, digester, path_stats.clone())
        .wait()
        .unwrap(),
      Snapshot {
        fingerprint: fingerprint,
        digest: Some(Digest(fingerprint, 78)),
        path_stats: path_stats,
      }
    );
  }

  #[test]
  fn snapshot_recursive_directories_including_empty() {
    let (store, dir, posix_fs, digester) = setup();

    let cats = PathBuf::from("cats");
    let roland = cats.join("roland");
    let dogs = PathBuf::from("dogs");
    let llamas = PathBuf::from("llamas");
    std::fs::create_dir_all(&dir.path().join(&cats)).unwrap();
    std::fs::create_dir_all(&dir.path().join(&dogs)).unwrap();
    std::fs::create_dir_all(&dir.path().join(&llamas)).unwrap();
    make_file(&dir.path().join(&roland), STR.as_bytes(), 0o600);

    let sorted_path_stats = expand_all_sorted(posix_fs);
    let mut unsorted_path_stats = sorted_path_stats.clone();
    unsorted_path_stats.reverse();
    // TODO: Inline when only used once
    let fingerprint = Fingerprint::from_hex_string(
      "fbff703bdaac62accf2ea5083bcfed89292073bf710ef9ad14d9298c637e777b",
    ).unwrap();
    assert_eq!(
      Snapshot::from_path_stats(store, digester, unsorted_path_stats)
        .wait()
        .unwrap(),
      Snapshot {
        fingerprint: fingerprint,
        digest: Some(Digest(fingerprint, 232)),
        path_stats: sorted_path_stats,
      }
    );
  }

  #[derive(Clone)]
  struct FileSaver(Arc<Store>, Arc<PosixFS>);

  impl StoreFileByDigest<String> for FileSaver {
    fn store_by_digest(&self, file: &File) -> BoxFuture<Digest, String> {
      let file_copy = file.clone();
      let store = self.0.clone();
      self
        .1
        .clone()
        .read_file(&file)
        .map_err(move |err| {
          format!("Error reading file {:?}: {}", file_copy, err.description())
        })
        .and_then(move |content| store.store_file_bytes(content.content, true))
        .to_boxed()
    }
  }

  fn expand_all_sorted(posix_fs: Arc<PosixFS>) -> Vec<PathStat> {
    let mut v = posix_fs
      .expand(PathGlobs::create(&["**".to_owned()], &vec![]).unwrap())
      .wait()
      .unwrap();
    v.sort_by(|a, b| a.path().cmp(b.path()));
    v
  }
}
<commit_msg>Snapshot can get FileContent (#5494)<commit_after>// Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use bazel_protos;
use boxfuture::{Boxable, BoxFuture};
use bytes::Bytes;
use futures::Future;
use futures::future::join_all;
use hashing::{Digest, Fingerprint};
use itertools::Itertools;
use {File, FileContent, PathStat, Store};
use protobuf;
use std::collections::HashMap;
use std::ffi::OsString;
use std::fmt;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

#[derive(Clone, PartialEq)]
pub struct Snapshot {
  // TODO: In a follow-up commit, fingerprint will be removed, and digest will be made non-optional.
  // They both exist right now as a compatibility shim so that the tar-based code and
  // Directory-based code can peacefully co-exist.
  pub fingerprint: Fingerprint,
  pub digest: Option<Digest>,
  pub path_stats: Vec<PathStat>,
}

// StoreFileByDigest allows a File to be saved to an underlying Store, in such a way that it can be
// looked up by the Digest produced by the store_by_digest method.
// It is a separate trait so that caching implementations can be written which wrap the Store (used
// to store the bytes) and VFS (used to read the files off disk if needed).
pub trait StoreFileByDigest<Error> {
  fn store_by_digest(&self, file: &File) -> BoxFuture<Digest, Error>;
}

impl Snapshot {
  pub fn from_path_stats<
    S: StoreFileByDigest<Error> + Sized + Clone,
    Error: fmt::Debug + 'static + Send,
  >(
    store: Arc<Store>,
    file_digester: S,
    path_stats: Vec<PathStat>,
  ) -> BoxFuture<Snapshot, String> {
    let mut sorted_path_stats = path_stats.clone();
    sorted_path_stats.sort_by(|a, b| a.path().cmp(b.path()));
    Snapshot::ingest_directory_from_sorted_path_stats(store, file_digester, sorted_path_stats)
      .map(|digest| {
        Snapshot {
          fingerprint: digest.0,
          digest: Some(digest),
          path_stats: path_stats,
        }
      })
      .to_boxed()
  }

  fn ingest_directory_from_sorted_path_stats<
    S: StoreFileByDigest<Error> + Sized + Clone,
    Error: fmt::Debug + 'static + Send,
  >(
    store: Arc<Store>,
    file_digester: S,
    path_stats: Vec<PathStat>,
  ) -> BoxFuture<Digest, String> {
    let mut file_futures: Vec<BoxFuture<bazel_protos::remote_execution::FileNode, String>> =
      Vec::new();
    let mut dir_futures: Vec<BoxFuture<bazel_protos::remote_execution::DirectoryNode, String>> =
      Vec::new();

    for (first_component, group) in
      &path_stats.iter().cloned().group_by(|s| {
        s.path().components().next().unwrap().as_os_str().to_owned()
      })
    {
      let mut path_group: Vec<PathStat> = group.collect();
      if path_group.len() == 1 && path_group.get(0).unwrap().path().components().count() == 1 {
        // Exactly one entry with exactly one component indicates either a file in this directory,
        // or an empty directory.
        // If the child is a non-empty directory, or a file therein, there must be multiple
        // PathStats with that prefix component, and we will handle that in the recursive
        // save_directory call.

        match path_group.pop().unwrap() {
          PathStat::File { ref stat, .. } => {
            let is_executable = stat.is_executable;
            file_futures.push(
              file_digester
                .clone()
                .store_by_digest(&stat)
                .map_err(|e| format!("{:?}", e))
                .and_then(move |digest| {
                  let mut file_node = bazel_protos::remote_execution::FileNode::new();
                  file_node.set_name(osstring_as_utf8(first_component)?);
                  file_node.set_digest((&digest).into());
                  file_node.set_is_executable(is_executable);
                  Ok(file_node)
                })
                .to_boxed(),
            );
          }
          PathStat::Dir { .. } => {
            // Because there are no children of this Dir, it must be empty.
            dir_futures.push(
              store
                .record_directory(&bazel_protos::remote_execution::Directory::new(), true)
                .map(move |digest| {
                  let mut directory_node = bazel_protos::remote_execution::DirectoryNode::new();
                  directory_node.set_name(osstring_as_utf8(first_component).unwrap());
                  directory_node.set_digest((&digest).into());
                  directory_node
                })
                .to_boxed(),
            );
          }
        }
      } else {
        dir_futures.push(
          // TODO: Memoize this in the graph
          Snapshot::ingest_directory_from_sorted_path_stats(
            store.clone(),
            file_digester.clone(),
            paths_of_child_dir(path_group),
          ).and_then(move |digest| {
            let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
            dir_node.set_name(osstring_as_utf8(first_component)?);
            dir_node.set_digest((&digest).into());
            Ok(dir_node)
          })
            .to_boxed(),
        );
      }
    }
    join_all(dir_futures)
      .join(join_all(file_futures))
      .and_then(move |(dirs, files)| {
        let mut directory = bazel_protos::remote_execution::Directory::new();
        directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
        directory.set_files(protobuf::RepeatedField::from_vec(files));
        store.record_directory(&directory, true)
      })
      .to_boxed()
  }

  // Preserves the order of Snapshot's path_stats in its returned Vec.
  pub fn contents(self, store: Arc<Store>) -> BoxFuture<Vec<FileContent>, String> {
    let contents = Arc::new(Mutex::new(HashMap::new()));
    let path_stats = self.path_stats;
    Snapshot::contents_for_directory_helper(
      self.digest.unwrap(),
      store,
      PathBuf::from(""),
      contents.clone(),
    ).map(move |_| {
      let mut contents = contents.lock().unwrap();
      let mut vec = Vec::new();
      for path in path_stats.iter().filter_map(|path_stat| match path_stat {
        &PathStat::File { ref path, .. } => Some(path.to_path_buf()),
        &PathStat::Dir { .. } => None,
      })
      {
        match contents.remove(&path) {
          Some(content) => vec.push(FileContent { path, content }),
          None => {
            panic!(format!(
              "PathStat for {:?} was present in path_stats but missing from Snapshot contents",
              path
            ));
          }
        }
      }
      vec
    })
      .to_boxed()
  }

  // Assumes that all fingerprints it encounters are valid.
  fn contents_for_directory_helper(
    digest: Digest,
    store: Arc<Store>,
    path_so_far: PathBuf,
    contents_wrapped: Arc<Mutex<HashMap<PathBuf, Bytes>>>,
  ) -> BoxFuture<(), String> {
    store
      .load_directory(digest)
      .and_then(move |maybe_dir| {
        maybe_dir.ok_or_else(|| {
          format!("Could not find directory with digest {:?}", digest)
        })
      })
      .and_then(move |dir| {
        let contents_wrapped_copy = contents_wrapped.clone();
        let path_so_far_copy = path_so_far.clone();
        let store_copy = store.clone();
        let file_futures = join_all(
          dir
            .get_files()
            .iter()
            .map(move |file_node| {
              let path = path_so_far_copy.join(file_node.get_name());
              let contents_wrapped_copy = contents_wrapped_copy.clone();
              store_copy
                .load_file_bytes_with(file_node.get_digest().into(), |b| b)
                .and_then(move |maybe_bytes| {
                  maybe_bytes
                    .ok_or_else(|| format!("Couldn't find file contents for {:?}", path))
                    .map(move |bytes| {
                      let mut contents = contents_wrapped_copy.lock().unwrap();
                      contents.insert(path, bytes);
                    })
                })
            })
            .collect::<Vec<_>>(),
        );
        let contents_wrapped_copy2 = contents_wrapped.clone();
        let store_copy = store.clone();
        let dir_futures = join_all(
          dir
            .get_directories()
            .iter()
            .map(move |dir_node| {
              Snapshot::contents_for_directory_helper(
                dir_node.get_digest().into(),
                store_copy.clone(),
                path_so_far.join(dir_node.get_name()),
                contents_wrapped_copy2.clone(),
              )
            })
            .collect::<Vec<_>>(),
        );
        file_futures.join(dir_futures)
      })
      .map(|(_, _)| ())
      .to_boxed()
  }
}

impl fmt::Debug for Snapshot {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(
      f,
      "Snapshot({}, digest={:?}, entries={})",
      self.fingerprint.to_hex(),
      self.digest,
      self.path_stats.len()
    )
  }
}

fn paths_of_child_dir(paths: Vec<PathStat>) -> Vec<PathStat> {
  paths
    .into_iter()
    .filter_map(|s| {
      if s.path().components().count() == 1 {
        return None;
      }
      Some(match s {
        PathStat::File { path, stat } => {
          PathStat::File {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
        PathStat::Dir { path, stat } => {
          PathStat::Dir {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
      })
    })
    .collect()
}

fn osstring_as_utf8(path: OsString) -> Result<String, String> {
  path.into_string().map_err(|p| {
    format!("{:?}'s file_name is not representable in UTF8", p)
  })
}

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

  use boxfuture::{BoxFuture, Boxable};
  use bytes::Bytes;
  use futures::future::Future;
  use hashing::{Digest, Fingerprint};
  use tempdir::TempDir;
  use self::testutil::make_file;

  use super::super::{File, FileContent, Path, PathGlobs, PathStat, PosixFS, ResettablePool,
                     Snapshot, Store, StoreFileByDigest, VFS};

  use std;
  use std::error::Error;
  use std::path::PathBuf;
  use std::sync::Arc;

  const AGGRESSIVE: &str = "Aggressive";
  const LATIN: &str = "Chaetophractus villosus";
  const STR: &str = "European Burmese";

  fn setup() -> (Arc<Store>, TempDir, Arc<PosixFS>, FileSaver) {
    let pool = Arc::new(ResettablePool::new("test-pool-".to_string()));
    // TODO: Pass a remote CAS address through.
    let store = Arc::new(
      Store::local_only(TempDir::new("lmdb_store").unwrap(), pool.clone()).unwrap(),
    );
    let dir = TempDir::new("root").unwrap();
    let posix_fs = Arc::new(PosixFS::new(dir.path(), pool, vec![]).unwrap());
    let file_saver = FileSaver(store.clone(), posix_fs.clone());
    (store, dir, posix_fs, file_saver)
  }

  #[test]
  fn snapshot_one_file() {
    let (store, dir, posix_fs, digester) = setup();

    let file_name = PathBuf::from("roland");
    make_file(&dir.path().join(&file_name), STR.as_bytes(), 0o600);

    let path_stats = expand_all_sorted(posix_fs);
    // TODO: Inline when only used once
    let fingerprint = Fingerprint::from_hex_string(
      "63949aa823baf765eff07b946050d76ec0033144c785a94d3ebd82baa931cd16",
    ).unwrap();
    assert_eq!(
      Snapshot::from_path_stats(store, digester, path_stats.clone())
        .wait()
        .unwrap(),
      Snapshot {
        fingerprint: fingerprint,
        digest: Some(Digest(fingerprint, 80)),
        path_stats: path_stats,
      }
    );
  }

  #[test]
  fn snapshot_recursive_directories() {
    let (store, dir, posix_fs, digester) = setup();

    let cats = PathBuf::from("cats");
    let roland = cats.join("roland");
    std::fs::create_dir_all(&dir.path().join(cats)).unwrap();
    make_file(&dir.path().join(&roland), STR.as_bytes(), 0o600);

    let path_stats = expand_all_sorted(posix_fs);
    // TODO: Inline when only used once
    let fingerprint = Fingerprint::from_hex_string(
      "8b1a7ea04eaa2527b35683edac088bc826117b53b7ec6601740b55e20bce3deb",
    ).unwrap();
    assert_eq!(
      Snapshot::from_path_stats(store, digester, path_stats.clone())
        .wait()
        .unwrap(),
      Snapshot {
        fingerprint: fingerprint,
        digest: Some(Digest(fingerprint, 78)),
        path_stats: path_stats,
      }
    );
  }

  #[test]
  fn snapshot_recursive_directories_including_empty() {
    let (store, dir, posix_fs, digester) = setup();

    let cats = PathBuf::from("cats");
    let roland = cats.join("roland");
    let dogs = PathBuf::from("dogs");
    let llamas = PathBuf::from("llamas");
    std::fs::create_dir_all(&dir.path().join(&cats)).unwrap();
    std::fs::create_dir_all(&dir.path().join(&dogs)).unwrap();
    std::fs::create_dir_all(&dir.path().join(&llamas)).unwrap();
    make_file(&dir.path().join(&roland), STR.as_bytes(), 0o600);

    let sorted_path_stats = expand_all_sorted(posix_fs);
    let mut unsorted_path_stats = sorted_path_stats.clone();
    unsorted_path_stats.reverse();
    // TODO: Inline when only used once
    let fingerprint = Fingerprint::from_hex_string(
      "fbff703bdaac62accf2ea5083bcfed89292073bf710ef9ad14d9298c637e777b",
    ).unwrap();
    assert_eq!(
      Snapshot::from_path_stats(store, digester, unsorted_path_stats.clone())
        .wait()
        .unwrap(),
      Snapshot {
        fingerprint: fingerprint,
        digest: Some(Digest(fingerprint, 232)),
        path_stats: unsorted_path_stats,
      }
    );
  }

  #[test]
  fn contents_for_one_file() {
    let (store, dir, posix_fs, digester) = setup();

    let file_name = PathBuf::from("roland");
    make_file(&dir.path().join(&file_name), STR.as_bytes(), 0o600);

    let contents = Snapshot::from_path_stats(store.clone(), digester, expand_all_sorted(posix_fs))
      .wait()
      .unwrap()
      .contents(store)
      .wait()
      .unwrap();
    assert_snapshot_contents(contents, vec![(&file_name, STR)]);
  }

  #[test]
  fn contents_for_files_in_multiple_directories() {
    let (store, dir, posix_fs, digester) = setup();

    let armadillos = PathBuf::from("armadillos");
    let armadillos_abs = dir.path().join(&armadillos);
    std::fs::create_dir_all(&armadillos_abs).unwrap();
    let amy = armadillos.join("amy");
    make_file(&dir.path().join(&amy), LATIN.as_bytes(), 0o600);
    let rolex = armadillos.join("rolex");
    make_file(&dir.path().join(&rolex), AGGRESSIVE.as_bytes(), 0o600);

    let cats = PathBuf::from("cats");
    let cats_abs = dir.path().join(&cats);
    std::fs::create_dir_all(&cats_abs).unwrap();
    let roland = cats.join("roland");
    make_file(&dir.path().join(&roland), STR.as_bytes(), 0o600);

    let dogs = PathBuf::from("dogs");
    let dogs_abs = dir.path().join(&dogs);
    std::fs::create_dir_all(&dogs_abs).unwrap();

    let path_stats_sorted = expand_all_sorted(posix_fs);
    let mut path_stats_reversed = path_stats_sorted.clone();
    path_stats_reversed.reverse();
    {
      let snapshot =
        Snapshot::from_path_stats(store.clone(), digester.clone(), path_stats_reversed)
          .wait()
          .unwrap();
      let contents = snapshot.contents(store.clone()).wait().unwrap();
      assert_snapshot_contents(
        contents,
        vec![(&roland, STR), (&rolex, AGGRESSIVE), (&amy, LATIN)],
      );
    }
    {
      let contents = Snapshot::from_path_stats(store.clone(), digester, path_stats_sorted)
        .wait()
        .unwrap()
        .contents(store)
        .wait()
        .unwrap();
      assert_snapshot_contents(
        contents,
        vec![(&amy, LATIN), (&rolex, AGGRESSIVE), (&roland, STR)],
      );
    }
  }

  #[derive(Clone)]
  struct FileSaver(Arc<Store>, Arc<PosixFS>);

  impl StoreFileByDigest<String> for FileSaver {
    fn store_by_digest(&self, file: &File) -> BoxFuture<Digest, String> {
      let file_copy = file.clone();
      let store = self.0.clone();
      self
        .1
        .clone()
        .read_file(&file)
        .map_err(move |err| {
          format!("Error reading file {:?}: {}", file_copy, err.description())
        })
        .and_then(move |content| store.store_file_bytes(content.content, true))
        .to_boxed()
    }
  }

  fn expand_all_sorted(posix_fs: Arc<PosixFS>) -> Vec<PathStat> {
    let mut v = posix_fs
      .expand(PathGlobs::create(&["**".to_owned()], &vec![]).unwrap())
      .wait()
      .unwrap();
    v.sort_by(|a, b| a.path().cmp(b.path()));
    v
  }

  fn assert_snapshot_contents(contents: Vec<FileContent>, expected: Vec<(&Path, &str)>) {
    let expected_with_array: Vec<_> = expected
      .into_iter()
      .map(|(path, s)| (path.to_path_buf(), Bytes::from(s)))
      .collect();
    let got: Vec<_> = contents
      .into_iter()
      .map(|file_content| (file_content.path, file_content.content))
      .collect();
    assert_eq!(expected_with_array, got);
  }
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use std::cmp;
use std::mem;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};

use glium::glutin::EventsLoopProxy;

use model::ttable::{Score, TTable};
use model::{Board, Move, Outcome};

const NEG_INFINITY: i16 = -0x7000;
const LOSE: i16 = -0x4000;
// Small contempt factor to discourage draws
const DRAW: i16 = 1;
const ASPIRATION_WIDTH: i16 = 51;

pub struct AI {
    status: Status,
    ttable: Arc<Mutex<TTable>>,
}

enum Status {
    Idle,
    // Either the AI thread is running, or there is a move waiting to be received
    Thinking {
        move_recv: Receiver<Move>,
        // We store and load this atomic with Ordering::Relaxed. It *should* be fine because it
        // doesn't interact with any other atomics--all we want to do is tell the AI thread to stop
        // searching for a move
        stop_signal: Arc<AtomicBool>,
        handle: JoinHandle<()>,
    },
}

impl AI {
    pub fn new() -> Self {
        Self {
            status: Status::Idle,
            ttable: Arc::new(Mutex::new(TTable::new())),
        }
    }

    pub fn is_idle(&self) -> bool {
        match self.status {
            Status::Idle => true,
            Status::Thinking { .. } => false,
        }
    }

    pub fn stop(&mut self) {
        if let Status::Thinking {
            ref stop_signal, ..
        } = self.status
        {
            stop_signal.store(true, Ordering::Relaxed);
        }
        // Unconditionally assign because without NLL, we can't put this in the if let block above
        self.status = Status::Idle;
    }

    pub fn try_recv(&mut self) -> Option<Move> {
        use self::TryRecvError::*;

        let result;
        match self.status {
            Status::Idle => result = None,
            Status::Thinking { ref move_recv, .. } => match move_recv.try_recv() {
                Ok(mv) => result = Some(mv),
                Err(Empty) => result = None,
                Err(Disconnected) => panic!("Tried to receive move from disconnected sender"),
            },
        }

        // We can't set status in the Ok(mv) arm above without NLL
        if result.is_some() {
            self.status = Status::Idle;
        }
        result
    }

    pub fn think(
        &mut self,
        board: Board,
        board_list: Vec<Board>,
        depth: u8,
        events_proxy: EventsLoopProxy,
    ) {
        assert_ne!(depth, 0);

        let prev_status = mem::replace(&mut self.status, Status::Idle);

        let (move_sender, move_recv) = mpsc::channel();
        let stop_signal = Arc::new(AtomicBool::new(false));
        let stop_signal_clone = stop_signal.clone();

        let ttable_mutex = self.ttable.clone();

        let handle = thread::spawn(move || {
            if let Status::Thinking {
                stop_signal,
                handle,
                ..
            } = prev_status
            {
                stop_signal.store(true, Ordering::Relaxed);
                handle
                    .join()
                    .expect("Old AI thread panicked when new AI thread joined on it");
            }

            // If the previous AI thread was sent the stop signal, but hasn't received it yet, we
            // will block here until it finishes. We won't have joined on its handle above because
            // stop throws away its Status enum.
            let mut ttable = match ttable_mutex.lock() {
                Ok(table) => table,
                Err(_) => panic!("Transposition table mutex is poisoned"),
            };

            if let SearchResult::Move(mv) =
                search_root(depth, board, board_list, &mut ttable, &stop_signal_clone)
            {
                if stop_signal_clone.load(Ordering::Relaxed) {
                    return;
                }
                move_sender.send(mv).expect("AI failed to send Move");
                events_proxy
                    .wakeup()
                    .expect("Failed to wake up events loop");
            }
        });

        self.status = Status::Thinking {
            move_recv,
            stop_signal,
            handle,
        };
    }
}

enum SearchResult {
    Move(Move),
    Stopped,
}

fn search_root(
    depth: u8,
    board: Board,
    board_list: Vec<Board>,
    ttable: &mut TTable,
    stop_signal: &Arc<AtomicBool>,
) -> SearchResult {
    ttable.inc_age();

    // Only take positions after the last irreversible move
    let mut board_list: Vec<_> = board_list
        .into_iter()
        .rev()
        .take_while(|b| b.vitals == board.vitals)
        .collect();
    board_list.reverse();

    let mut moves: Vec<(Move, i16)> = board
        .generate_moves()
        .map(|mv| (mv, NEG_INFINITY))
        .collect();

    if moves.is_empty() {
        panic!("AI has no moves");
    }

    let mut pv = None;
    let mut iter_score = evaluate(&board);
    for depth in 0..depth {
        if stop_signal.load(Ordering::Relaxed) {
            return SearchResult::Stopped;
        }

        // Aspiration window search loop
        let mut asp_width = ASPIRATION_WIDTH;
        loop {
            let mut max_score = iter_score - asp_width;
            for pair in &mut moves {
                let mut new_board = board;
                new_board.apply_move(&pair.0);

                let mut new_pv = vec![];

                let score = -alphabeta_negamax(
                    &new_board,
                    &mut board_list,
                    &mut new_pv,
                    -(iter_score + asp_width),
                    -max_score,
                    depth,
                    ttable,
                );

                if score > max_score {
                    max_score = score;
                    pv = Some(new_pv);
                }
                pair.1 = score;
            }

            if max_score == iter_score + asp_width || max_score == iter_score - asp_width {
                // True score lies outside the window, so we search this position again
                asp_width *= 2;
            } else {
                break;
            }
        }

        moves.sort_by(|&(_, a), &(_, b)| b.cmp(&a));
        iter_score = moves[0].1;

        println!("\nDepth {}: {:>6}", depth, moves[0].1);
        if let Some(ref mut pv) = pv {
            pv.push(moves[0].0);
            for mv in pv.iter().rev() {
                println!("    {}", mv);
            }
        }
    }
    println!("\n---------------------");

    SearchResult::Move(moves[0].0)
}

fn alphabeta_negamax(
    board: &Board,
    // This list does not include the current board
    mut board_list: &mut Vec<Board>,
    pv: &mut Vec<Move>,
    mut alpha: i16,
    mut beta: i16,
    depth: u8,
    ttable: &mut TTable,
) -> i16 {
    let mut set_pv = move |score, new_pv| {
        if score > alpha && score < beta {
            *pv = new_pv;
        }
    };
    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth as i8);
    };

    use self::Outcome::*;
    match board.outcome() {
        DrawStalemate | DrawInsufficientMaterial => {
            // This is safe to do because Board does not detect draws by threefold repetition
            set_ttable(ttable, Score::Exact(DRAW));
            set_pv(DRAW, vec![]);
            return DRAW;
        }
        Win(color) => {
            assert_ne!(color, board.turn);
            // Weight score by depth to encourage shorter wins. The shorter the win, the greater
            // `depth` will be, and so the larger the score will be. This also encourages the AI to
            // prolong a loss.
            let score = LOSE - i16::from(depth);
            set_ttable(ttable, Score::Exact(score));
            set_pv(score, vec![]);
            return score;
        }
        InProgress => {}
        DrawThreefoldRepetition => unreachable!(),
    }

    if board_list.len() >= 8 && board_list.iter().filter(|&&b| b == *board).count() >= 2 {
        set_pv(DRAW, vec![]);
        return DRAW;
    }

    if depth == 0 {
        let score = quiescence_search(board, alpha, beta, depth as i8, ttable);
        set_pv(score, vec![]);
        return score;
    }

    {
        match ttable.get(board.zobrist, depth as i8) {
            Some(Score::Exact(score)) => {
                // This will cut the PV short
                // TODO: Store the best move in the table and get the PV from that?
                set_pv(score, vec![]);
                return score;
            }
            Some(Score::Beta(score)) => {
                if score >= beta {
                    return score;
                }
                beta = score;
            }
            None => {}
        }
    }

    let mut best_score = NEG_INFINITY;
    let mut best_move = None;

    let mut new_pv = vec![];
    for mv in board.generate_moves() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        board_list.push(*board);
        let score = -alphabeta_negamax(
            &new_board,
            &mut board_list,
            &mut new_pv,
            -beta,
            -alpha,
            depth - 1,
            ttable,
        );
        board_list.pop();

        best_score = cmp::max(score, best_score);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
            best_move = Some(mv);
        }
    }
    set_ttable(ttable, Score::Exact(best_score));
    if let Some(mv) = best_move {
        new_pv.push(mv);
        set_pv(alpha, new_pv);
    }
    alpha
}

fn quiescence_search(
    board: &Board,
    mut alpha: i16,
    mut beta: i16,
    depth: i8,
    ttable: &mut TTable,
) -> i16 {
    let stand_pat = evaluate(board);
    if stand_pat >= beta {
        return beta;
    } else if stand_pat + 200 < alpha {
        // Delta pruning: we bet that no capture move will raise alpha by more than 200 centipieces
        // We don't switch this off in the endgame because a draw by insufficient material is very
        // unlikely (both sides must be unable to exchange and each have only 1 piece left)
        return alpha;
    } else if alpha < stand_pat {
        alpha = stand_pat;
    }

    match ttable.get(board.zobrist, depth) {
        Some(Score::Exact(score)) => {
            return score;
        }
        Some(Score::Beta(score)) => {
            if score >= beta {
                return score;
            }
            beta = score;
        }
        None => {}
    };

    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth);
    };

    for mv in board.generate_captures() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        let score = -quiescence_search(&new_board, -beta, -alpha, depth - 1, ttable);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
        }
    }
    set_ttable(ttable, Score::Exact(alpha));
    alpha
}

fn evaluate(board: &Board) -> i16 {
    use model::Color::*;

    let wp = 100 * i16::from(board.pieces(White));
    let bp = 100 * i16::from(board.pieces(Black));
    let wh = 50 * i16::from(board.hexes(White));
    let bh = 50 * i16::from(board.hexes(Black));

    match board.turn {
        White => (wp + wh) - (bp + bh),
        Black => (bp + bh) - (wp + wh),
    }
}
<commit_msg>Use error name to appease Clippy<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use std::cmp;
use std::mem;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};

use glium::glutin::EventsLoopProxy;

use model::ttable::{Score, TTable};
use model::{Board, Move, Outcome};

const NEG_INFINITY: i16 = -0x7000;
const LOSE: i16 = -0x4000;
// Small contempt factor to discourage draws
const DRAW: i16 = 1;
const ASPIRATION_WIDTH: i16 = 51;

pub struct AI {
    status: Status,
    ttable: Arc<Mutex<TTable>>,
}

enum Status {
    Idle,
    // Either the AI thread is running, or there is a move waiting to be received
    Thinking {
        move_recv: Receiver<Move>,
        // We store and load this atomic with Ordering::Relaxed. It *should* be fine because it
        // doesn't interact with any other atomics--all we want to do is tell the AI thread to stop
        // searching for a move
        stop_signal: Arc<AtomicBool>,
        handle: JoinHandle<()>,
    },
}

impl AI {
    pub fn new() -> Self {
        Self {
            status: Status::Idle,
            ttable: Arc::new(Mutex::new(TTable::new())),
        }
    }

    pub fn is_idle(&self) -> bool {
        match self.status {
            Status::Idle => true,
            Status::Thinking { .. } => false,
        }
    }

    pub fn stop(&mut self) {
        if let Status::Thinking {
            ref stop_signal, ..
        } = self.status
        {
            stop_signal.store(true, Ordering::Relaxed);
        }
        // Unconditionally assign because without NLL, we can't put this in the if let block above
        self.status = Status::Idle;
    }

    pub fn try_recv(&mut self) -> Option<Move> {
        use self::TryRecvError::*;

        let result;
        match self.status {
            Status::Idle => result = None,
            Status::Thinking { ref move_recv, .. } => match move_recv.try_recv() {
                Ok(mv) => result = Some(mv),
                Err(Empty) => result = None,
                Err(Disconnected) => panic!("Tried to receive move from disconnected sender"),
            },
        }

        // We can't set status in the Ok(mv) arm above without NLL
        if result.is_some() {
            self.status = Status::Idle;
        }
        result
    }

    pub fn think(
        &mut self,
        board: Board,
        board_list: Vec<Board>,
        depth: u8,
        events_proxy: EventsLoopProxy,
    ) {
        assert_ne!(depth, 0);

        let prev_status = mem::replace(&mut self.status, Status::Idle);

        let (move_sender, move_recv) = mpsc::channel();
        let stop_signal = Arc::new(AtomicBool::new(false));
        let stop_signal_clone = stop_signal.clone();

        let ttable_mutex = self.ttable.clone();

        let handle = thread::spawn(move || {
            if let Status::Thinking {
                stop_signal,
                handle,
                ..
            } = prev_status
            {
                stop_signal.store(true, Ordering::Relaxed);
                handle
                    .join()
                    .expect("Old AI thread panicked when new AI thread joined on it");
            }

            // If the previous AI thread was sent the stop signal, but hasn't received it yet, we
            // will block here until it finishes. We won't have joined on its handle above because
            // stop throws away its Status enum.
            let mut ttable = match ttable_mutex.lock() {
                Ok(table) => table,
                Err(_poison_error) => panic!("Transposition table mutex is poisoned"),
            };

            if let SearchResult::Move(mv) =
                search_root(depth, board, board_list, &mut ttable, &stop_signal_clone)
            {
                if stop_signal_clone.load(Ordering::Relaxed) {
                    return;
                }
                move_sender.send(mv).expect("AI failed to send Move");
                events_proxy
                    .wakeup()
                    .expect("Failed to wake up events loop");
            }
        });

        self.status = Status::Thinking {
            move_recv,
            stop_signal,
            handle,
        };
    }
}

enum SearchResult {
    Move(Move),
    Stopped,
}

fn search_root(
    depth: u8,
    board: Board,
    board_list: Vec<Board>,
    ttable: &mut TTable,
    stop_signal: &Arc<AtomicBool>,
) -> SearchResult {
    ttable.inc_age();

    // Only take positions after the last irreversible move
    let mut board_list: Vec<_> = board_list
        .into_iter()
        .rev()
        .take_while(|b| b.vitals == board.vitals)
        .collect();
    board_list.reverse();

    let mut moves: Vec<(Move, i16)> = board
        .generate_moves()
        .map(|mv| (mv, NEG_INFINITY))
        .collect();

    if moves.is_empty() {
        panic!("AI has no moves");
    }

    let mut pv = None;
    let mut iter_score = evaluate(&board);
    for depth in 0..depth {
        if stop_signal.load(Ordering::Relaxed) {
            return SearchResult::Stopped;
        }

        // Aspiration window search loop
        let mut asp_width = ASPIRATION_WIDTH;
        loop {
            let mut max_score = iter_score - asp_width;
            for pair in &mut moves {
                let mut new_board = board;
                new_board.apply_move(&pair.0);

                let mut new_pv = vec![];

                let score = -alphabeta_negamax(
                    &new_board,
                    &mut board_list,
                    &mut new_pv,
                    -(iter_score + asp_width),
                    -max_score,
                    depth,
                    ttable,
                );

                if score > max_score {
                    max_score = score;
                    pv = Some(new_pv);
                }
                pair.1 = score;
            }

            if max_score == iter_score + asp_width || max_score == iter_score - asp_width {
                // True score lies outside the window, so we search this position again
                asp_width *= 2;
            } else {
                break;
            }
        }

        moves.sort_by(|&(_, a), &(_, b)| b.cmp(&a));
        iter_score = moves[0].1;

        println!("\nDepth {}: {:>6}", depth, moves[0].1);
        if let Some(ref mut pv) = pv {
            pv.push(moves[0].0);
            for mv in pv.iter().rev() {
                println!("    {}", mv);
            }
        }
    }
    println!("\n---------------------");

    SearchResult::Move(moves[0].0)
}

fn alphabeta_negamax(
    board: &Board,
    // This list does not include the current board
    mut board_list: &mut Vec<Board>,
    pv: &mut Vec<Move>,
    mut alpha: i16,
    mut beta: i16,
    depth: u8,
    ttable: &mut TTable,
) -> i16 {
    let mut set_pv = move |score, new_pv| {
        if score > alpha && score < beta {
            *pv = new_pv;
        }
    };
    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth as i8);
    };

    use self::Outcome::*;
    match board.outcome() {
        DrawStalemate | DrawInsufficientMaterial => {
            // This is safe to do because Board does not detect draws by threefold repetition
            set_ttable(ttable, Score::Exact(DRAW));
            set_pv(DRAW, vec![]);
            return DRAW;
        }
        Win(color) => {
            assert_ne!(color, board.turn);
            // Weight score by depth to encourage shorter wins. The shorter the win, the greater
            // `depth` will be, and so the larger the score will be. This also encourages the AI to
            // prolong a loss.
            let score = LOSE - i16::from(depth);
            set_ttable(ttable, Score::Exact(score));
            set_pv(score, vec![]);
            return score;
        }
        InProgress => {}
        DrawThreefoldRepetition => unreachable!(),
    }

    if board_list.len() >= 8 && board_list.iter().filter(|&&b| b == *board).count() >= 2 {
        set_pv(DRAW, vec![]);
        return DRAW;
    }

    if depth == 0 {
        let score = quiescence_search(board, alpha, beta, depth as i8, ttable);
        set_pv(score, vec![]);
        return score;
    }

    {
        match ttable.get(board.zobrist, depth as i8) {
            Some(Score::Exact(score)) => {
                // This will cut the PV short
                // TODO: Store the best move in the table and get the PV from that?
                set_pv(score, vec![]);
                return score;
            }
            Some(Score::Beta(score)) => {
                if score >= beta {
                    return score;
                }
                beta = score;
            }
            None => {}
        }
    }

    let mut best_score = NEG_INFINITY;
    let mut best_move = None;

    let mut new_pv = vec![];
    for mv in board.generate_moves() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        board_list.push(*board);
        let score = -alphabeta_negamax(
            &new_board,
            &mut board_list,
            &mut new_pv,
            -beta,
            -alpha,
            depth - 1,
            ttable,
        );
        board_list.pop();

        best_score = cmp::max(score, best_score);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
            best_move = Some(mv);
        }
    }
    set_ttable(ttable, Score::Exact(best_score));
    if let Some(mv) = best_move {
        new_pv.push(mv);
        set_pv(alpha, new_pv);
    }
    alpha
}

fn quiescence_search(
    board: &Board,
    mut alpha: i16,
    mut beta: i16,
    depth: i8,
    ttable: &mut TTable,
) -> i16 {
    let stand_pat = evaluate(board);
    if stand_pat >= beta {
        return beta;
    } else if stand_pat + 200 < alpha {
        // Delta pruning: we bet that no capture move will raise alpha by more than 200 centipieces
        // We don't switch this off in the endgame because a draw by insufficient material is very
        // unlikely (both sides must be unable to exchange and each have only 1 piece left)
        return alpha;
    } else if alpha < stand_pat {
        alpha = stand_pat;
    }

    match ttable.get(board.zobrist, depth) {
        Some(Score::Exact(score)) => {
            return score;
        }
        Some(Score::Beta(score)) => {
            if score >= beta {
                return score;
            }
            beta = score;
        }
        None => {}
    };

    let set_ttable = |ttable: &mut TTable, score| {
        ttable.set(board.zobrist, score, depth);
    };

    for mv in board.generate_captures() {
        let mut new_board = *board;
        new_board.apply_move(&mv);

        let score = -quiescence_search(&new_board, -beta, -alpha, depth - 1, ttable);

        if score >= beta {
            set_ttable(ttable, Score::Beta(score));
            return beta;
        } else if score > alpha {
            alpha = score;
        }
    }
    set_ttable(ttable, Score::Exact(alpha));
    alpha
}

fn evaluate(board: &Board) -> i16 {
    use model::Color::*;

    let wp = 100 * i16::from(board.pieces(White));
    let bp = 100 * i16::from(board.pieces(Black));
    let wh = 50 * i16::from(board.hexes(White));
    let bh = 50 * i16::from(board.hexes(Black));

    match board.turn {
        White => (wp + wh) - (bp + bh),
        Black => (bp + bh) - (wp + wh),
    }
}
<|endoftext|>


<commit_before>use clap::ArgMatches;
use std::env;
use std::process::Command;

pub fn open_shell_at_repo(args: &ArgMatches) -> Result<(), (String, i32)> {
    info!("CD to: {}", args.value_of("repository").unwrap());

    let shell = match env::var("SHELL") {
        Ok(val) => val,
        Err(e) => return Err((format!("Could not determine shell to use: {}", e), 1)),
    };

    info!("Using shell: {}", shell);
    match Command::new(shell).status() {
        Ok(exit_status) => {
            if exit_status.success() {
                Ok(())
            } else {
                Err(("Shell failed".to_string(), exit_status.code().unwrap_or(1)))
            }
        },
        Err(e) => {
            Err((format!("Unable to execute shell: {}", e), 1))
        }
    }
}
<commit_msg>Provide more context when shell command runs, but fails<commit_after>use clap::ArgMatches;
use std::env;
use std::process::Command;

pub fn open_shell_at_repo(args: &ArgMatches) -> Result<(), (String, i32)> {
    info!("CD to: {}", args.value_of("repository").unwrap());

    let shell = match env::var("SHELL") {
        Ok(val) => val,
        Err(e) => return Err((format!("Could not determine shell to use: {}", e), 1)),
    };

    info!("Using shell: {}", shell);
    match Command::new(shell).status() {
        Ok(exit_status) => {
            if exit_status.success() {
                Ok(())
            } else {
                let exit_code = exit_status.code().unwrap_or(1);
                Err((format!("Shell failed (exit {})", exit_code), exit_code))
            }
        },
        Err(e) => {
            Err((format!("Unable to execute shell: {}", e), 1))
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Hash Function
//!
//! Primary hash function used in the protocol
//!

use byteorder::{BigEndian, ByteOrder};
use std::cmp::min;
use std::convert::AsRef;
use std::ops::Add;
use std::{fmt, ops};

use crate::blake2::blake2b::Blake2b;

use crate::ser::{self, AsFixedBytes, Error, FixedLength, Readable, Reader, Writeable, Writer};
use crate::util;

/// A hash consisting of all zeroes, used as a sentinel. No known preimage.
pub const ZERO_HASH: Hash = Hash([0; 32]);

/// A hash to uniquely (or close enough) identify one of the main blockchain
/// constructs. Used pervasively for blocks, transactions and outputs.
#[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Serialize, Deserialize)]
pub struct Hash([u8; 32]);

impl fmt::Debug for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		for i in self.0[..4].iter() {
			write!(f, "{:02x}", i)?;
		}
		Ok(())
	}
}

impl fmt::Display for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let hash_hex = self.to_hex();
		let len = hash_hex.len();
		const NUM_SHOW: usize = 8;

		let starting = &hash_hex[..NUM_SHOW];
		let ending = &hash_hex[(len - NUM_SHOW)..];

		write!(f, "{}...{}", starting, ending)
	}
}

impl FixedLength for Hash {
	/// Size of a hash in bytes.
	const LEN: usize = 32;
}

impl Hash {
	/// Builds a Hash from a byte vector. If the vector is too short, it will be
	/// completed by zeroes. If it's too long, it will be truncated.
	pub fn from_vec(v: &[u8]) -> Hash {
		let mut h = [0; Hash::LEN];
		let copy_size = min(v.len(), Hash::LEN);
		h[..copy_size].copy_from_slice(&v[..copy_size]);
		Hash(h)
	}

	/// Converts the hash to a byte vector
	pub fn to_vec(&self) -> Vec<u8> {
		self.0.to_vec()
	}

	/// Returns a byte slice of the hash contents.
	pub fn as_bytes(&self) -> &[u8] {
		&self.0
	}

	/// Convert a hash to hex string format.
	pub fn to_hex(&self) -> String {
		util::to_hex(self.to_vec())
	}

	/// Convert hex string back to hash.
	pub fn from_hex(hex: &str) -> Result<Hash, Error> {
		let bytes = util::from_hex(hex.to_string())
			.map_err(|_| Error::HexError(format!("failed to decode {}", hex)))?;
		Ok(Hash::from_vec(&bytes))
	}

	/// Most significant 64 bits
	pub fn to_u64(&self) -> u64 {
		BigEndian::read_u64(&self.0)
	}
}

impl ops::Index<usize> for Hash {
	type Output = u8;

	fn index(&self, idx: usize) -> &u8 {
		&self.0[idx]
	}
}

impl ops::Index<ops::Range<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::Range<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeTo<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeTo<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFrom<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFrom<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFull> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFull) -> &[u8] {
		&self.0[idx]
	}
}

impl AsRef<[u8]> for Hash {
	fn as_ref(&self) -> &[u8] {
		&self.0
	}
}

impl Readable for Hash {
	fn read(reader: &mut dyn Reader) -> Result<Hash, ser::Error> {
		let v = reader.read_fixed_bytes(32)?;
		let mut a = [0; 32];
		a.copy_from_slice(&v[..]);
		Ok(Hash(a))
	}
}

impl Writeable for Hash {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), Error> {
		writer.write_fixed_bytes(&self.0)
	}
}

impl Add for Hash {
	type Output = Hash;
	fn add(self, other: Hash) -> Hash {
		self.hash_with(other)
	}
}

impl Default for Hash {
	fn default() -> Hash {
		ZERO_HASH
	}
}

/// Serializer that outputs a hash of the serialized object
pub struct HashWriter {
	state: Blake2b,
}

impl HashWriter {
	/// Consume the `HashWriter`, outputting its current hash into a 32-byte
	/// array
	pub fn finalize(self, output: &mut [u8]) {
		output.copy_from_slice(self.state.finalize().as_bytes());
	}

	/// Consume the `HashWriter`, outputting a `Hash` corresponding to its
	/// current state
	pub fn into_hash(self) -> Hash {
		let mut res = [0; 32];
		res.copy_from_slice(self.state.finalize().as_bytes());
		Hash(res)
	}
}

impl Default for HashWriter {
	fn default() -> HashWriter {
		HashWriter {
			state: Blake2b::new(32),
		}
	}
}

impl ser::Writer for HashWriter {
	fn serialization_mode(&self) -> ser::SerializationMode {
		ser::SerializationMode::Hash
	}

	fn write_fixed_bytes<T: AsFixedBytes>(&mut self, b32: &T) -> Result<(), ser::Error> {
		self.state.update(b32.as_ref());
		Ok(())
	}
}

/// A trait for types that have a canonical hash
pub trait Hashed {
	/// Obtain the hash of the object
	fn hash(&self) -> Hash;
	/// Hash the object together with another writeable object
	fn hash_with<T: Writeable>(&self, other: T) -> Hash;
}

impl<W: ser::Writeable> Hashed for W {
	fn hash(&self) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}

	fn hash_with<T: Writeable>(&self, other: T) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		ser::Writeable::write(&other, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}
}
<commit_msg>Show first 12 digits of Hash<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Hash Function
//!
//! Primary hash function used in the protocol
//!

use byteorder::{BigEndian, ByteOrder};
use std::cmp::min;
use std::convert::AsRef;
use std::ops::Add;
use std::{fmt, ops};

use crate::blake2::blake2b::Blake2b;

use crate::ser::{self, AsFixedBytes, Error, FixedLength, Readable, Reader, Writeable, Writer};
use crate::util;

/// A hash consisting of all zeroes, used as a sentinel. No known preimage.
pub const ZERO_HASH: Hash = Hash([0; 32]);

/// A hash to uniquely (or close enough) identify one of the main blockchain
/// constructs. Used pervasively for blocks, transactions and outputs.
#[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Serialize, Deserialize)]
pub struct Hash([u8; 32]);

impl fmt::Debug for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		for i in self.0[..4].iter() {
			write!(f, "{:02x}", i)?;
		}
		Ok(())
	}
}

impl fmt::Display for Hash {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let hash_hex = self.to_hex();
		const NUM_SHOW: usize = 12;

		write!(f, "{}...", &hash_hex[..NUM_SHOW])
	}
}

impl FixedLength for Hash {
	/// Size of a hash in bytes.
	const LEN: usize = 32;
}

impl Hash {
	/// Builds a Hash from a byte vector. If the vector is too short, it will be
	/// completed by zeroes. If it's too long, it will be truncated.
	pub fn from_vec(v: &[u8]) -> Hash {
		let mut h = [0; Hash::LEN];
		let copy_size = min(v.len(), Hash::LEN);
		h[..copy_size].copy_from_slice(&v[..copy_size]);
		Hash(h)
	}

	/// Converts the hash to a byte vector
	pub fn to_vec(&self) -> Vec<u8> {
		self.0.to_vec()
	}

	/// Returns a byte slice of the hash contents.
	pub fn as_bytes(&self) -> &[u8] {
		&self.0
	}

	/// Convert a hash to hex string format.
	pub fn to_hex(&self) -> String {
		util::to_hex(self.to_vec())
	}

	/// Convert hex string back to hash.
	pub fn from_hex(hex: &str) -> Result<Hash, Error> {
		let bytes = util::from_hex(hex.to_string())
			.map_err(|_| Error::HexError(format!("failed to decode {}", hex)))?;
		Ok(Hash::from_vec(&bytes))
	}

	/// Most significant 64 bits
	pub fn to_u64(&self) -> u64 {
		BigEndian::read_u64(&self.0)
	}
}

impl ops::Index<usize> for Hash {
	type Output = u8;

	fn index(&self, idx: usize) -> &u8 {
		&self.0[idx]
	}
}

impl ops::Index<ops::Range<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::Range<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeTo<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeTo<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFrom<usize>> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFrom<usize>) -> &[u8] {
		&self.0[idx]
	}
}

impl ops::Index<ops::RangeFull> for Hash {
	type Output = [u8];

	fn index(&self, idx: ops::RangeFull) -> &[u8] {
		&self.0[idx]
	}
}

impl AsRef<[u8]> for Hash {
	fn as_ref(&self) -> &[u8] {
		&self.0
	}
}

impl Readable for Hash {
	fn read(reader: &mut dyn Reader) -> Result<Hash, ser::Error> {
		let v = reader.read_fixed_bytes(32)?;
		let mut a = [0; 32];
		a.copy_from_slice(&v[..]);
		Ok(Hash(a))
	}
}

impl Writeable for Hash {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), Error> {
		writer.write_fixed_bytes(&self.0)
	}
}

impl Add for Hash {
	type Output = Hash;
	fn add(self, other: Hash) -> Hash {
		self.hash_with(other)
	}
}

impl Default for Hash {
	fn default() -> Hash {
		ZERO_HASH
	}
}

/// Serializer that outputs a hash of the serialized object
pub struct HashWriter {
	state: Blake2b,
}

impl HashWriter {
	/// Consume the `HashWriter`, outputting its current hash into a 32-byte
	/// array
	pub fn finalize(self, output: &mut [u8]) {
		output.copy_from_slice(self.state.finalize().as_bytes());
	}

	/// Consume the `HashWriter`, outputting a `Hash` corresponding to its
	/// current state
	pub fn into_hash(self) -> Hash {
		let mut res = [0; 32];
		res.copy_from_slice(self.state.finalize().as_bytes());
		Hash(res)
	}
}

impl Default for HashWriter {
	fn default() -> HashWriter {
		HashWriter {
			state: Blake2b::new(32),
		}
	}
}

impl ser::Writer for HashWriter {
	fn serialization_mode(&self) -> ser::SerializationMode {
		ser::SerializationMode::Hash
	}

	fn write_fixed_bytes<T: AsFixedBytes>(&mut self, b32: &T) -> Result<(), ser::Error> {
		self.state.update(b32.as_ref());
		Ok(())
	}
}

/// A trait for types that have a canonical hash
pub trait Hashed {
	/// Obtain the hash of the object
	fn hash(&self) -> Hash;
	/// Hash the object together with another writeable object
	fn hash_with<T: Writeable>(&self, other: T) -> Hash;
}

impl<W: ser::Writeable> Hashed for W {
	fn hash(&self) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}

	fn hash_with<T: Writeable>(&self, other: T) -> Hash {
		let mut hasher = HashWriter::default();
		ser::Writeable::write(self, &mut hasher).unwrap();
		ser::Writeable::write(&other, &mut hasher).unwrap();
		let mut ret = [0; 32];
		hasher.finalize(&mut ret);
		Hash(ret)
	}
}
<|endoftext|>


<commit_before>// https://leetcode.com/problems/largest-merge-of-two-strings/
pub fn largest_merge(word1: String, word2: String) -> String {
    todo!()
}

fn main() {
    println!(
        "{}",
        largest_merge(String::from("cabaa"), String::from("bcaaa"))
    ); // cbcabaaaaa
    println!(
        "{}",
        largest_merge(String::from("abcabc"), String::from("abdcaba"))
    ); // abdcabcabcaba
}
<commit_msg>Update problem 1754<commit_after>// https://leetcode.com/problems/largest-merge-of-two-strings/
pub fn largest_merge(word1: String, word2: String) -> String {
    todo!()
}

fn main() {
    println!(
        "{}",
        largest_merge(String::from("cabaa"), String::from("bcaaa"))
    ); // "cbcabaaaaa"
    println!(
        "{}",
        largest_merge(String::from("abcabc"), String::from("abdcaba"))
    ); // "abdcabcabcaba"
    println!(
        "{}",
        largest_merge(String::from("abb"), String::from("acde"))
    ); // "acdeabb"
}
<|endoftext|>