<commit_before>//! The JIT driver uses [`cranelift_jit`] to JIT execute programs without writing any object
//! files.

use std::cell::RefCell;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
use rustc_codegen_ssa::CrateInfo;
use rustc_middle::mir::mono::MonoItem;

use cranelift_jit::{JITBuilder, JITModule};

use crate::{prelude::*, BackendConfig};
use crate::{CodegenCx, CodegenMode};

struct JitState {
    backend_config: BackendConfig,
    jit_module: JITModule,
}

thread_local! {
    static LAZY_JIT_STATE: RefCell<Option<JitState>> = RefCell::new(None);
}

fn create_jit_module<'tcx>(
    tcx: TyCtxt<'tcx>,
    backend_config: &BackendConfig,
    hotswap: bool,
) -> (JITModule, CodegenCx<'tcx>) {
    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let isa = crate::build_isa(tcx.sess, backend_config);
    let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
    jit_builder.hotswap(hotswap);
    crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
    jit_builder.symbols(imported_symbols);
    let mut jit_module = JITModule::new(jit_builder);

    let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);

    crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
    crate::main_shim::maybe_create_entry_wrapper(
        tcx,
        &mut jit_module,
        &mut cx.unwind_context,
        true,
        true,
    );

    (jit_module, cx)
}

pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
    if !tcx.sess.opts.output_types.should_codegen() {
        tcx.sess.fatal("JIT mode doesn't work with `cargo check`");
    }

    if !tcx.sess.crate_types().contains(&rustc_session::config::CrateType::Executable) {
        tcx.sess.fatal("can't jit non-executable crate");
    }

    let (mut jit_module, mut cx) = create_jit_module(
        tcx,
        &backend_config,
        matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
    );

    let (_, cgus) = tcx.collect_and_partition_mono_items(());
    let mono_items = cgus
        .iter()
        .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
        .flatten()
        .collect::<FxHashMap<_, (_, _)>>()
        .into_iter()
        .collect::<Vec<(_, (_, _))>>();

    super::time(tcx, backend_config.display_cg_time, "codegen mono items", || {
        super::predefine_mono_items(tcx, &mut jit_module, &mono_items);
        for (mono_item, _) in mono_items {
            match mono_item {
                MonoItem::Fn(inst) => match backend_config.codegen_mode {
                    CodegenMode::Aot => unreachable!(),
                    CodegenMode::Jit => {
                        cx.tcx.sess.time("codegen fn", || {
                            crate::base::codegen_fn(&mut cx, &mut jit_module, inst)
                        });
                    }
                    CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
                },
                MonoItem::Static(def_id) => {
                    crate::constant::codegen_static(tcx, &mut jit_module, def_id);
                }
                MonoItem::GlobalAsm(item_id) => {
                    let item = tcx.hir().item(item_id);
                    tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
                }
            }
        }
    });

    if !cx.global_asm.is_empty() {
        tcx.sess.fatal("Inline asm is not supported in JIT mode");
    }

    tcx.sess.abort_if_errors();

    jit_module.finalize_definitions();
    unsafe { cx.unwind_context.register_jit(&jit_module) };

    println!(
        "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
    );

    let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
        .chain(backend_config.jit_args.iter().map(|arg| &**arg))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();

    // Push a null pointer as a terminating argument. This is required by POSIX and
    // useful as some dynamic linkers use it as a marker to jump over.
    argv.push(std::ptr::null());

    let start_sig = Signature {
        params: vec![
            AbiParam::new(jit_module.target_config().pointer_type()),
            AbiParam::new(jit_module.target_config().pointer_type()),
        ],
        returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
        call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
    };
    let start_func_id = jit_module.declare_function("main", Linkage::Import, &start_sig).unwrap();
    let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);

    LAZY_JIT_STATE.with(|lazy_jit_state| {
        let mut lazy_jit_state = lazy_jit_state.borrow_mut();
        assert!(lazy_jit_state.is_none());
        *lazy_jit_state = Some(JitState { backend_config, jit_module });
    });

    let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
        unsafe { ::std::mem::transmute(finalized_start) };
    let ret = f(args.len() as c_int, argv.as_ptr());
    std::process::exit(ret);
}

#[no_mangle]
extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
    rustc_middle::ty::tls::with(|tcx| {
        // lift is used to ensure the correct lifetime for instance.
        let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();

        LAZY_JIT_STATE.with(|lazy_jit_state| {
            let mut lazy_jit_state = lazy_jit_state.borrow_mut();
            let lazy_jit_state = lazy_jit_state.as_mut().unwrap();
            let jit_module = &mut lazy_jit_state.jit_module;
            let backend_config = lazy_jit_state.backend_config.clone();

            let name = tcx.symbol_name(instance).name;
            let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
            let func_id = jit_module.declare_function(name, Linkage::Export, &sig).unwrap();
            jit_module.prepare_for_function_redefine(func_id).unwrap();

            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module.isa(), false);
            tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, jit_module, instance));

            assert!(cx.global_asm.is_empty());
            jit_module.finalize_definitions();
            unsafe { cx.unwind_context.register_jit(&jit_module) };
            jit_module.get_finalized_function(func_id)
        })
    })
}

fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
    use rustc_middle::middle::dependency_format::Linkage;

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(());
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
        .unwrap()
        .1;
    for &(cnum, _) in &crate_info.used_crates_dynamic {
        let src = &crate_info.used_crate_source[&cnum];
        match data[cnum.as_usize() - 1] {
            Linkage::NotLinked | Linkage::IncludedFromDylib => {}
            Linkage::Static => {
                let name = tcx.crate_name(cnum);
                let mut err =
                    tcx.sess.struct_err(&format!("Can't load static lib {}", name.as_str()));
                err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
                err.emit();
            }
            Linkage::Dynamic => {
                dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
            }
        }
    }

    let mut imported_symbols = Vec::new();
    for path in dylib_paths {
        use object::{Object, ObjectSymbol};
        let lib = libloading::Library::new(&path).unwrap();
        let obj = std::fs::read(path).unwrap();
        let obj = object::File::parse(&*obj).unwrap();
        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                return None;
            }
            if name.starts_with("rust_metadata_") {
                // The metadata is part of a section that is not loaded by the dynamic linker in
                // case of cg_llvm.
                return None;
            }
            let dlsym_name = if cfg!(target_os = "macos") {
                // On macOS `dlsym` expects the name without leading `_`.
                assert!(name.starts_with('_'), "{:?}", name);
                &name[1..]
            } else {
                &name
            };
            let symbol: libloading::Symbol<'_, *const u8> =
                unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
            Some((name, *symbol))
        }));
        std::mem::forget(lib)
    }

    tcx.sess.abort_if_errors();

    imported_symbols
}

fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
    let tcx = cx.tcx;

    let pointer_type = module.target_config().pointer_type();

    let name = tcx.symbol_name(inst).name;
    let sig = crate::abi::get_function_sig(tcx, module.isa().triple(), inst);
    let func_id = module.declare_function(name, Linkage::Export, &sig).unwrap();

    let instance_ptr = Box::into_raw(Box::new(inst));

    let jit_fn = module
        .declare_function(
            "__clif_jit_fn",
            Linkage::Import,
            &Signature {
                call_conv: module.target_config().default_call_conv,
                params: vec![AbiParam::new(pointer_type)],
                returns: vec![AbiParam::new(pointer_type)],
            },
        )
        .unwrap();

    cx.cached_context.clear();
    let trampoline = &mut cx.cached_context.func;
    trampoline.signature = sig.clone();

    let mut builder_ctx = FunctionBuilderContext::new();
    let mut trampoline_builder = FunctionBuilder::new(trampoline, &mut builder_ctx);

    let jit_fn = module.declare_func_in_func(jit_fn, trampoline_builder.func);
    let sig_ref = trampoline_builder.func.import_signature(sig);

    let entry_block = trampoline_builder.create_block();
    trampoline_builder.append_block_params_for_function_params(entry_block);
    let fn_args = trampoline_builder.func.dfg.block_params(entry_block).to_vec();

    trampoline_builder.switch_to_block(entry_block);
    let instance_ptr = trampoline_builder.ins().iconst(pointer_type, instance_ptr as u64 as i64);
    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
    let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
    let call_inst = trampoline_builder.ins().call_indirect(sig_ref, jitted_fn, &fn_args);
    let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
    trampoline_builder.ins().return_(&ret_vals);

    module
        .define_function(
            func_id,
            &mut cx.cached_context,
            &mut NullTrapSink {},
            &mut NullStackMapSink {},
        )
        .unwrap();
}
<commit_msg>Remove LibSource<commit_after>//! The JIT driver uses [`cranelift_jit`] to JIT execute programs without writing any object
//! files.

use std::cell::RefCell;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
use rustc_codegen_ssa::CrateInfo;
use rustc_middle::mir::mono::MonoItem;

use cranelift_jit::{JITBuilder, JITModule};

use crate::{prelude::*, BackendConfig};
use crate::{CodegenCx, CodegenMode};

struct JitState {
    backend_config: BackendConfig,
    jit_module: JITModule,
}

thread_local! {
    static LAZY_JIT_STATE: RefCell<Option<JitState>> = RefCell::new(None);
}

fn create_jit_module<'tcx>(
    tcx: TyCtxt<'tcx>,
    backend_config: &BackendConfig,
    hotswap: bool,
) -> (JITModule, CodegenCx<'tcx>) {
    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let isa = crate::build_isa(tcx.sess, backend_config);
    let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
    jit_builder.hotswap(hotswap);
    crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
    jit_builder.symbols(imported_symbols);
    let mut jit_module = JITModule::new(jit_builder);

    let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);

    crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
    crate::main_shim::maybe_create_entry_wrapper(
        tcx,
        &mut jit_module,
        &mut cx.unwind_context,
        true,
        true,
    );

    (jit_module, cx)
}

pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
    if !tcx.sess.opts.output_types.should_codegen() {
        tcx.sess.fatal("JIT mode doesn't work with `cargo check`");
    }

    if !tcx.sess.crate_types().contains(&rustc_session::config::CrateType::Executable) {
        tcx.sess.fatal("can't jit non-executable crate");
    }

    let (mut jit_module, mut cx) = create_jit_module(
        tcx,
        &backend_config,
        matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
    );

    let (_, cgus) = tcx.collect_and_partition_mono_items(());
    let mono_items = cgus
        .iter()
        .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
        .flatten()
        .collect::<FxHashMap<_, (_, _)>>()
        .into_iter()
        .collect::<Vec<(_, (_, _))>>();

    super::time(tcx, backend_config.display_cg_time, "codegen mono items", || {
        super::predefine_mono_items(tcx, &mut jit_module, &mono_items);
        for (mono_item, _) in mono_items {
            match mono_item {
                MonoItem::Fn(inst) => match backend_config.codegen_mode {
                    CodegenMode::Aot => unreachable!(),
                    CodegenMode::Jit => {
                        cx.tcx.sess.time("codegen fn", || {
                            crate::base::codegen_fn(&mut cx, &mut jit_module, inst)
                        });
                    }
                    CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
                },
                MonoItem::Static(def_id) => {
                    crate::constant::codegen_static(tcx, &mut jit_module, def_id);
                }
                MonoItem::GlobalAsm(item_id) => {
                    let item = tcx.hir().item(item_id);
                    tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
                }
            }
        }
    });

    if !cx.global_asm.is_empty() {
        tcx.sess.fatal("Inline asm is not supported in JIT mode");
    }

    tcx.sess.abort_if_errors();

    jit_module.finalize_definitions();
    unsafe { cx.unwind_context.register_jit(&jit_module) };

    println!(
        "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
    );

    let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
        .chain(backend_config.jit_args.iter().map(|arg| &**arg))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();

    // Push a null pointer as a terminating argument. This is required by POSIX and
    // useful as some dynamic linkers use it as a marker to jump over.
    argv.push(std::ptr::null());

    let start_sig = Signature {
        params: vec![
            AbiParam::new(jit_module.target_config().pointer_type()),
            AbiParam::new(jit_module.target_config().pointer_type()),
        ],
        returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
        call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
    };
    let start_func_id = jit_module.declare_function("main", Linkage::Import, &start_sig).unwrap();
    let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);

    LAZY_JIT_STATE.with(|lazy_jit_state| {
        let mut lazy_jit_state = lazy_jit_state.borrow_mut();
        assert!(lazy_jit_state.is_none());
        *lazy_jit_state = Some(JitState { backend_config, jit_module });
    });

    let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
        unsafe { ::std::mem::transmute(finalized_start) };
    let ret = f(args.len() as c_int, argv.as_ptr());
    std::process::exit(ret);
}

#[no_mangle]
extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
    rustc_middle::ty::tls::with(|tcx| {
        // lift is used to ensure the correct lifetime for instance.
        let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();

        LAZY_JIT_STATE.with(|lazy_jit_state| {
            let mut lazy_jit_state = lazy_jit_state.borrow_mut();
            let lazy_jit_state = lazy_jit_state.as_mut().unwrap();
            let jit_module = &mut lazy_jit_state.jit_module;
            let backend_config = lazy_jit_state.backend_config.clone();

            let name = tcx.symbol_name(instance).name;
            let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
            let func_id = jit_module.declare_function(name, Linkage::Export, &sig).unwrap();
            jit_module.prepare_for_function_redefine(func_id).unwrap();

            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module.isa(), false);
            tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, jit_module, instance));

            assert!(cx.global_asm.is_empty());
            jit_module.finalize_definitions();
            unsafe { cx.unwind_context.register_jit(&jit_module) };
            jit_module.get_finalized_function(func_id)
        })
    })
}

fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
    use rustc_middle::middle::dependency_format::Linkage;

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(());
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
        .unwrap()
        .1;
    for &cnum in &crate_info.used_crates {
        let src = &crate_info.used_crate_source[&cnum];
        match data[cnum.as_usize() - 1] {
            Linkage::NotLinked | Linkage::IncludedFromDylib => {}
            Linkage::Static => {
                let name = tcx.crate_name(cnum);
                let mut err =
                    tcx.sess.struct_err(&format!("Can't load static lib {}", name.as_str()));
                err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
                err.emit();
            }
            Linkage::Dynamic => {
                dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
            }
        }
    }

    let mut imported_symbols = Vec::new();
    for path in dylib_paths {
        use object::{Object, ObjectSymbol};
        let lib = libloading::Library::new(&path).unwrap();
        let obj = std::fs::read(path).unwrap();
        let obj = object::File::parse(&*obj).unwrap();
        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                return None;
            }
            if name.starts_with("rust_metadata_") {
                // The metadata is part of a section that is not loaded by the dynamic linker in
                // case of cg_llvm.
                return None;
            }
            let dlsym_name = if cfg!(target_os = "macos") {
                // On macOS `dlsym` expects the name without leading `_`.
                assert!(name.starts_with('_'), "{:?}", name);
                &name[1..]
            } else {
                &name
            };
            let symbol: libloading::Symbol<'_, *const u8> =
                unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
            Some((name, *symbol))
        }));
        std::mem::forget(lib)
    }

    tcx.sess.abort_if_errors();

    imported_symbols
}

fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
    let tcx = cx.tcx;

    let pointer_type = module.target_config().pointer_type();

    let name = tcx.symbol_name(inst).name;
    let sig = crate::abi::get_function_sig(tcx, module.isa().triple(), inst);
    let func_id = module.declare_function(name, Linkage::Export, &sig).unwrap();

    let instance_ptr = Box::into_raw(Box::new(inst));

    let jit_fn = module
        .declare_function(
            "__clif_jit_fn",
            Linkage::Import,
            &Signature {
                call_conv: module.target_config().default_call_conv,
                params: vec![AbiParam::new(pointer_type)],
                returns: vec![AbiParam::new(pointer_type)],
            },
        )
        .unwrap();

    cx.cached_context.clear();
    let trampoline = &mut cx.cached_context.func;
    trampoline.signature = sig.clone();

    let mut builder_ctx = FunctionBuilderContext::new();
    let mut trampoline_builder = FunctionBuilder::new(trampoline, &mut builder_ctx);

    let jit_fn = module.declare_func_in_func(jit_fn, trampoline_builder.func);
    let sig_ref = trampoline_builder.func.import_signature(sig);

    let entry_block = trampoline_builder.create_block();
    trampoline_builder.append_block_params_for_function_params(entry_block);
    let fn_args = trampoline_builder.func.dfg.block_params(entry_block).to_vec();

    trampoline_builder.switch_to_block(entry_block);
    let instance_ptr = trampoline_builder.ins().iconst(pointer_type, instance_ptr as u64 as i64);
    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
    let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
    let call_inst = trampoline_builder.ins().call_indirect(sig_ref, jitted_fn, &fn_args);
    let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
    trampoline_builder.ins().return_(&ret_vals);

    module
        .define_function(
            func_id,
            &mut cx.cached_context,
            &mut NullTrapSink {},
            &mut NullStackMapSink {},
        )
        .unwrap();
}
<|endoftext|>


<commit_before>extern crate neat;
extern crate rand;
extern crate graph_neighbor_matching;
extern crate graph_io_gml;
extern crate closed01;
extern crate petgraph;
extern crate cppn;

mod common;

use neat::population::{Population, Unrated, Runner};
use neat::genomes::acyclic_network::{Genome, GenomeDistance, Environment, ElementStrategy};
use neat::fitness::Fitness;
use neat::crossover::ProbabilisticCrossover;
use graph_neighbor_matching::{SimilarityMatrix, ScoreNorm};
use graph_neighbor_matching::graph::{OwnedGraph, GraphBuilder};
use rand::{Rng, Closed01};
use std::marker::PhantomData;
use neat::mutate::MutateMethodWeighting;
use common::{load_graph, Mater, Neuron, NodeColors, convert_neuron_from_str};
use cppn::cppn::{Cppn, CppnNodeType};
use cppn::bipolar::BipolarActivationFunction;
use cppn::substrate::Substrate;
use cppn::position::Position2d;

type Node = CppnNodeType<BipolarActivationFunction>;

#[derive(Debug)]
struct FitnessEvaluator {
    target_graph: OwnedGraph<Neuron>,
}

impl FitnessEvaluator {
    // A larger fitness means "better"
    fn fitness(&self, genome: &Genome<Node>) -> f32 {
        let mut substrate = Substrate::new();
        substrate.add_node(Position2d::new(-1.0, -1.0), Neuron::Input);
        substrate.add_node(Position2d::new(1.0, -1.0), Neuron::Input);
        substrate.add_node(Position2d::new(-1.0, 1.0), Neuron::Output);
        substrate.add_node(Position2d::new(0.0, 1.0), Neuron::Output);
        substrate.add_node(Position2d::new(1.0, 1.0), Neuron::Output);

        let mut cppn = Cppn::new(genome.network());

        // now develop the cppn. the result is a graph
        let mut builder = GraphBuilder::new();
        for (i, node) in substrate.nodes().iter().enumerate() {
            let _ = builder.add_node(i, node.data.clone());
        }
        for link in substrate.iter_links(&mut cppn, None) {
            if link.weight >= 0.0 && link.weight <= 1.0 {
                builder.add_edge(link.source_idx,
                                 link.target_idx,
                                 closed01::Closed01::new(link.weight as f32));
            }
        }
        let graph = builder.graph();
        println!("graph: {:#?}", graph);

        let mut s = SimilarityMatrix::new(&graph, &self.target_graph, NodeColors);
        s.iterate(50, 0.01);
        s.score_optimal_sum_norm(None, ScoreNorm::MaxDegree).get()
    }
}

struct ES;

impl ElementStrategy<Node> for ES {
    fn random_link_weight<R: Rng>(rng: &mut R) -> f64 {
        // XXX Choose a weight between -1 and 1?
        rng.gen()
    }

    fn random_node_type<R: Rng>(rng: &mut R) -> Node {
        let af = &[BipolarActivationFunction::Identity,
                   BipolarActivationFunction::Linear,
                   BipolarActivationFunction::Gaussian,
                   BipolarActivationFunction::Sigmoid,
                   BipolarActivationFunction::Sine];

        CppnNodeType::Hidden(*rng.choose(af).unwrap())
    }
}

const POP_SIZE: usize = 100;

fn main() {
    let mut rng = rand::thread_rng();

    let fitness_evaluator = FitnessEvaluator {
        target_graph: load_graph("examples/jeffress.gml", convert_neuron_from_str),
    };

    println!("{:?}", fitness_evaluator);

    // start with minimal random topology.
    let mut env: Environment<Node, ES> = Environment::new();

    let template_genome = {
        let mut genome = Genome::new();

        // 4 inputs (x1,y1,x2,y2)
        for _ in 0..4 {
            env.add_node_to_genome(&mut genome, CppnNodeType::Input);
        }

        // 1 output (y)
        env.add_node_to_genome(&mut genome, CppnNodeType::Output);

        // 1 bias node
        env.add_node_to_genome(&mut genome, CppnNodeType::Bias);

        genome
    };

    println!("{:#?}", template_genome);

    let mut initial_pop = Population::<_, Unrated>::new();

    for _ in 0..POP_SIZE {
        initial_pop.add_genome(Box::new(template_genome.clone()));
    }
    assert!(initial_pop.len() == POP_SIZE);

    let mut mater = Mater {
        p_crossover: Closed01(0.5),
        p_crossover_detail: ProbabilisticCrossover {
            prob_match_left: Closed01(0.5), /* NEAT always selects a random parent for matching genes */
            prob_disjoint_left: Closed01(0.9),
            prob_excess_left: Closed01(0.9),
            prob_disjoint_right: Closed01(0.15),
            prob_excess_right: Closed01(0.15),
        },
        mutate_weights: MutateMethodWeighting {
            w_modify_weight: 1,
            w_add_node: 1,
            w_add_connection: 1,
        },

        env: &mut env,
    };

    let compatibility = GenomeDistance {
        excess: 1.0,
        disjoint: 1.0,
        weight: 0.0,
    };

    let mut runner = Runner {
        pop_size: POP_SIZE,
        elite_percentage: Closed01(0.05),
        selection_percentage: Closed01(0.2),
        tournament_k: 3,
        compatibility_threshold: 1.0,
        compatibility: &compatibility,
        mate: &mut mater,
        fitness: &|genome| Fitness::new(fitness_evaluator.fitness(genome) as f64),
        _marker: PhantomData,
    };

    let (iter, new_pop) = runner.run(initial_pop,
                                     &|iter, pop| {
                                         iter >= 100 || pop.max_fitness().unwrap().get() > 0.99
                                     },
                                     &mut rng);

    let new_pop = new_pop.sort();

    println!("iter: {}", iter);
    println!("{:#?}", new_pop);
}
<commit_msg>Comment out debug output<commit_after>extern crate neat;
extern crate rand;
extern crate graph_neighbor_matching;
extern crate graph_io_gml;
extern crate closed01;
extern crate petgraph;
extern crate cppn;

mod common;

use neat::population::{Population, Unrated, Runner};
use neat::genomes::acyclic_network::{Genome, GenomeDistance, Environment, ElementStrategy};
use neat::fitness::Fitness;
use neat::crossover::ProbabilisticCrossover;
use graph_neighbor_matching::{SimilarityMatrix, ScoreNorm};
use graph_neighbor_matching::graph::{OwnedGraph, GraphBuilder};
use rand::{Rng, Closed01};
use std::marker::PhantomData;
use neat::mutate::MutateMethodWeighting;
use common::{load_graph, Mater, Neuron, NodeColors, convert_neuron_from_str};
use cppn::cppn::{Cppn, CppnNodeType};
use cppn::bipolar::BipolarActivationFunction;
use cppn::substrate::Substrate;
use cppn::position::Position2d;

type Node = CppnNodeType<BipolarActivationFunction>;

#[derive(Debug)]
struct FitnessEvaluator {
    target_graph: OwnedGraph<Neuron>,
}

impl FitnessEvaluator {
    // A larger fitness means "better"
    fn fitness(&self, genome: &Genome<Node>) -> f32 {
        let mut substrate = Substrate::new();
        substrate.add_node(Position2d::new(-1.0, -1.0), Neuron::Input);
        substrate.add_node(Position2d::new(1.0, -1.0), Neuron::Input);
        substrate.add_node(Position2d::new(-1.0, 1.0), Neuron::Output);
        substrate.add_node(Position2d::new(0.0, 1.0), Neuron::Output);
        substrate.add_node(Position2d::new(1.0, 1.0), Neuron::Output);

        let mut cppn = Cppn::new(genome.network());

        // now develop the cppn. the result is a graph
        let mut builder = GraphBuilder::new();
        for (i, node) in substrate.nodes().iter().enumerate() {
            let _ = builder.add_node(i, node.data.clone());
        }
        for link in substrate.iter_links(&mut cppn, None) {
            if link.weight >= 0.0 && link.weight <= 1.0 {
                builder.add_edge(link.source_idx,
                                 link.target_idx,
                                 closed01::Closed01::new(link.weight as f32));
            }
        }
        let graph = builder.graph();
        //println!("graph: {:#?}", graph);

        let mut s = SimilarityMatrix::new(&graph, &self.target_graph, NodeColors);
        s.iterate(50, 0.01);
        s.score_optimal_sum_norm(None, ScoreNorm::MaxDegree).get()
    }
}

struct ES;

impl ElementStrategy<Node> for ES {
    fn random_link_weight<R: Rng>(rng: &mut R) -> f64 {
        // XXX Choose a weight between -1 and 1?
        rng.gen()
    }

    fn random_node_type<R: Rng>(rng: &mut R) -> Node {
        let af = &[BipolarActivationFunction::Identity,
                   BipolarActivationFunction::Linear,
                   BipolarActivationFunction::Gaussian,
                   BipolarActivationFunction::Sigmoid,
                   BipolarActivationFunction::Sine];

        CppnNodeType::Hidden(*rng.choose(af).unwrap())
    }
}

const POP_SIZE: usize = 100;

fn main() {
    let mut rng = rand::thread_rng();

    let fitness_evaluator = FitnessEvaluator {
        target_graph: load_graph("examples/jeffress.gml", convert_neuron_from_str),
    };

    println!("{:?}", fitness_evaluator);

    // start with minimal random topology.
    let mut env: Environment<Node, ES> = Environment::new();

    let template_genome = {
        let mut genome = Genome::new();

        // 4 inputs (x1,y1,x2,y2)
        for _ in 0..4 {
            env.add_node_to_genome(&mut genome, CppnNodeType::Input);
        }

        // 1 output (y)
        env.add_node_to_genome(&mut genome, CppnNodeType::Output);

        // 1 bias node
        env.add_node_to_genome(&mut genome, CppnNodeType::Bias);

        genome
    };

    println!("{:#?}", template_genome);

    let mut initial_pop = Population::<_, Unrated>::new();

    for _ in 0..POP_SIZE {
        initial_pop.add_genome(Box::new(template_genome.clone()));
    }
    assert!(initial_pop.len() == POP_SIZE);

    let mut mater = Mater {
        p_crossover: Closed01(0.5),
        p_crossover_detail: ProbabilisticCrossover {
            prob_match_left: Closed01(0.5), /* NEAT always selects a random parent for matching genes */
            prob_disjoint_left: Closed01(0.9),
            prob_excess_left: Closed01(0.9),
            prob_disjoint_right: Closed01(0.15),
            prob_excess_right: Closed01(0.15),
        },
        mutate_weights: MutateMethodWeighting {
            w_modify_weight: 1,
            w_add_node: 1,
            w_add_connection: 1,
        },

        env: &mut env,
    };

    let compatibility = GenomeDistance {
        excess: 1.0,
        disjoint: 1.0,
        weight: 0.0,
    };

    let mut runner = Runner {
        pop_size: POP_SIZE,
        elite_percentage: Closed01(0.05),
        selection_percentage: Closed01(0.2),
        tournament_k: 3,
        compatibility_threshold: 1.0,
        compatibility: &compatibility,
        mate: &mut mater,
        fitness: &|genome| Fitness::new(fitness_evaluator.fitness(genome) as f64),
        _marker: PhantomData,
    };

    let (iter, new_pop) = runner.run(initial_pop,
                                     &|iter, pop| {
                                         iter >= 100 || pop.max_fitness().unwrap().get() > 0.99
                                     },
                                     &mut rng);

    let new_pop = new_pop.sort();

    println!("iter: {}", iter);
    println!("{:#?}", new_pop);
}
<|endoftext|>


<commit_before>use futures::task::{ArcWake, Context, Poll, Waker};
use futures::Stream;
use parking_lot::Mutex;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::pin::Pin;
use std::sync::atomic;
use std::sync::Arc;

pub(crate) struct QueueInner<S, K> {
    counter: atomic::AtomicUsize,
    ready_queue: BinaryHeap<PriorityStream<S, K>>,
    pending_streams: HashMap<usize, PriorityStream<S, K>>,
    // See wake_by_ref for details
    weak_up_processing: Option<usize>,
    waker: Option<Waker>,
}

impl<S, K> QueueInner<S, K> {
    pub fn insert(&mut self, k: K, s: S) {
        self.ready_queue.push(PriorityStream {
            priority: self.counter.fetch_add(1, atomic::Ordering::Relaxed),
            key: k,
            stream: Box::pin(s),
        });
        match &self.waker {
            Some(w) => w.wake_by_ref(),
            None => (),
        };
    }
}

pub struct FairQueue<S, K> {
    block_on_no_clients: bool,
    inner: Arc<Mutex<QueueInner<S, K>>>,
}

struct StreamWaker<S, K> {
    inner: Arc<Mutex<QueueInner<S, K>>>,
    index: usize,
}

impl<S, K> ArcWake for StreamWaker<S, K>
where
    S: Send,
    K: Send,
{
    fn wake_by_ref(arc_self: &Arc<Self>) {
        let mut inner = arc_self.inner.lock();
        match inner.pending_streams.remove(&arc_self.index) {
            None => {
                // This is a tricky part..
                // Some streams call waker inside the poll_next method.
                // At that moment stream is neither ready or pending.
                // We leave it's priority hang for the moment.
                // It's responsibility of the FairQueue::poll_next to take this into account.
                // In such case it will put stream as ready (cause it explicitly asked for it)
                inner.weak_up_processing = Some(arc_self.index);
            }
            Some(s) => {
                inner.ready_queue.push(s);
            }
        };
        if let Some(waker) = inner.waker.take() {
            waker.wake_by_ref();
        }
    }
}

struct PriorityStream<S, K> {
    priority: usize,
    key: K,
    stream: Pin<Box<S>>,
}

impl<S, K> PartialEq for PriorityStream<S, K> {
    fn eq(&self, other: &Self) -> bool {
        self.priority.eq(&other.priority)
    }
}
impl<S, K> Eq for PriorityStream<S, K> {}

impl<S, K> PartialOrd for PriorityStream<S, K> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.priority.partial_cmp(&self.priority)
    }
}
impl<S, K> Ord for PriorityStream<S, K> {
    fn cmp(&self, other: &Self) -> Ordering {
        other.priority.cmp(&self.priority)
    }
}

impl<S, T, K> Stream for FairQueue<S, K>
where
    T: Send,
    S: Stream<Item = T> + Send,
    K: Unpin + Clone + Send,
{
    type Item = (K, T);

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let stream = self.get_mut();
        loop {
            let mut s = {
                let mut inner = stream.inner.lock();
                inner.waker = Some(cx.waker().clone());
                match inner.ready_queue.pop() {
                    Some(s) => s,
                    None => {
                        return if !inner.pending_streams.is_empty() || stream.block_on_no_clients {
                            Poll::Pending
                        } else {
                            Poll::Ready(None)
                        }
                    }
                }
            };

            let waker = Arc::new(StreamWaker {
                inner: stream.inner.clone(),
                index: s.priority,
            });
            let waker_ref = futures::task::waker_ref(&waker);
            let mut cx = Context::from_waker(&waker_ref);
            match s.stream.as_mut().poll_next(&mut cx) {
                Poll::Ready(Some(res)) => {
                    let item = Some((s.key.clone(), res));
                    let mut inner = stream.inner.lock();
                    s.priority = inner.counter.fetch_add(1, atomic::Ordering::Relaxed);
                    inner.ready_queue.push(s);
                    inner.weak_up_processing = None;
                    return Poll::Ready(item);
                }
                Poll::Ready(None) => continue,
                Poll::Pending => {
                    let mut inner = stream.inner.lock();
                    match inner.weak_up_processing.take() {
                        None => {
                            inner.pending_streams.insert(s.priority, s);
                        }
                        Some(prio) => {
                            assert_eq!(prio, s.priority);
                            inner.ready_queue.push(s);
                        }
                    };
                    return Poll::Pending;
                }
            }
        }
    }
}

impl<S, K> FairQueue<S, K> {
    pub fn new(block_on_no_clients: bool) -> Self {
        Self {
            block_on_no_clients,
            inner: Arc::new(Mutex::new(QueueInner {
                counter: atomic::AtomicUsize::new(0),
                ready_queue: BinaryHeap::new(),
                pending_streams: HashMap::new(),
                weak_up_processing: None,
                waker: None,
            })),
        }
    }

    pub(crate) fn inner(&self) -> Arc<Mutex<QueueInner<S, K>>> {
        self.inner.clone()
    }
}

#[cfg(test)]
mod test {
    use crate::fair_queue::FairQueue;
    use futures::StreamExt;

    #[tokio::test]
    async fn test_fair_queue_ready() {
        let a = futures::stream::iter(vec!["a1", "a2", "a3"]);
        let b = futures::stream::iter(vec!["b1", "b2", "b3"]);
        let c = futures::stream::iter(vec!["c1", "c2", "c3"]);

        let mut f_queue: FairQueue<_, u64> = FairQueue::new(false);
        {
            let inner = f_queue.inner();
            let mut inner_lock = inner.lock();
            inner_lock.insert(1, a);
            inner_lock.insert(2, b);
            inner_lock.insert(3, c);
        }

        let mut results = Vec::new();
        while let Some(i) = f_queue.next().await {
            results.push(i);
        }
        assert_eq!(
            results,
            vec![
                (1, "a1"),
                (2, "b1"),
                (3, "c1"),
                (1, "a2"),
                (2, "b2"),
                (3, "c2"),
                (1, "a3"),
                (2, "b3"),
                (3, "c3")
            ]
        );
    }

    #[tokio::test]
    async fn test_fair_queue_different_size() {
        let a = futures::stream::iter(vec!["a1", "a2", "a3"]);
        let b = futures::stream::iter(vec!["b1"]);
        let c = futures::stream::iter(vec!["c1", "c2"]);

        let mut f_queue: FairQueue<_, u64> = FairQueue::new(false);
        {
            let inner = f_queue.inner();
            let mut inner_lock = inner.lock();
            inner_lock.insert(1, a);
            inner_lock.insert(2, b);
            inner_lock.insert(3, c);
        }

        let mut results = Vec::new();
        while let Some(i) = f_queue.next().await {
            results.push(i);
        }
        assert_eq!(
            results,
            vec![
                (1, "a1"),
                (2, "b1"),
                (3, "c1"),
                (1, "a2"),
                (3, "c2"),
                (1, "a3")
            ]
        );
    }
}
<commit_msg>Fixes #101<commit_after>use futures::task::{ArcWake, Context, Poll, Waker};
use futures::Stream;
use parking_lot::Mutex;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::hash::Hash;
use std::pin::Pin;
use std::sync::atomic;
use std::sync::Arc;

pub(crate) struct QueueInner<S, K: Clone> {
    counter: atomic::AtomicUsize,
    ready_queue: BinaryHeap<ReadyEvent<K>>,
    streams: HashMap<K, Pin<Box<S>>>,
    waker: Option<Waker>,
}

impl<S, K: Clone + Eq + Hash> QueueInner<S, K> {
    pub fn insert(&mut self, k: K, s: S) {
        self.streams.insert(k.clone(), Box::pin(s));
        self.ready_queue.push(ReadyEvent {
            priority: self.counter.fetch_add(1, atomic::Ordering::Relaxed),
            key: k,
        });
        match &self.waker {
            Some(w) => w.wake_by_ref(),
            None => (),
        };
    }
}

pub struct FairQueue<S, K: Clone> {
    block_on_no_clients: bool,
    inner: Arc<Mutex<QueueInner<S, K>>>,
}

#[derive(Clone)]
struct ReadyEvent<K: Clone> {
    priority: usize,
    key: K,
}

impl<K: Clone> PartialEq for ReadyEvent<K> {
    fn eq(&self, other: &Self) -> bool {
        self.priority.eq(&other.priority)
    }
}
impl<K: Clone> Eq for ReadyEvent<K> {}

impl<K: Clone> PartialOrd for ReadyEvent<K> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.priority.partial_cmp(&self.priority)
    }
}
impl<K: Clone> Ord for ReadyEvent<K> {
    fn cmp(&self, other: &Self) -> Ordering {
        other.priority.cmp(&self.priority)
    }
}

struct StreamWaker<S, K: Clone> {
    inner: Arc<Mutex<QueueInner<S, K>>>,
    event: ReadyEvent<K>,
}

impl<S, K> ArcWake for StreamWaker<S, K>
where
    S: Send,
    K: Clone + Send + Sync,
{
    fn wake_by_ref(arc_self: &Arc<Self>) {
        let mut inner = arc_self.inner.lock();
        inner.ready_queue.push(arc_self.event.clone());
        if let Some(waker) = inner.waker.take() {
            waker.wake_by_ref();
        }
    }
}

impl<S, T, K> Stream for FairQueue<S, K>
where
    T: Send,
    S: Stream<Item = T> + Send,
    K: Eq + Hash + Unpin + Clone + Send + Sync,
{
    type Item = (K, T);

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let fair_queue = self.get_mut();
        loop {
            let (event, mut io_stream) = {
                let mut inner = fair_queue.inner.lock();
                inner.waker = Some(cx.waker().clone());
                let event = match inner.ready_queue.pop() {
                    Some(s) => s,
                    None => {
                        return if !inner.streams.is_empty() || fair_queue.block_on_no_clients {
                            Poll::Pending
                        } else {
                            Poll::Ready(None)
                        }
                    }
                };
                match inner.streams.remove(&event.key) {
                    Some(stream) => (event, stream),
                    None => continue,
                }
            };

            let waker = Arc::new(StreamWaker {
                inner: fair_queue.inner.clone(),
                event: event.clone(),
            });
            let waker_ref = futures::task::waker_ref(&waker);
            let mut cx = Context::from_waker(&waker_ref);
            match io_stream.as_mut().poll_next(&mut cx) {
                Poll::Ready(Some(res)) => {
                    let item = Some((event.key.clone(), res));
                    let mut inner = fair_queue.inner.lock();
                    let priority = inner.counter.fetch_add(1, atomic::Ordering::Relaxed);
                    inner.ready_queue.push(ReadyEvent {
                        priority,
                        key: event.key.clone(),
                    });
                    inner.streams.insert(event.key, io_stream);
                    return Poll::Ready(item);
                }
                Poll::Ready(None) => continue,
                Poll::Pending => {
                    let mut inner = fair_queue.inner.lock();
                    inner.streams.insert(event.key, io_stream);
                    return Poll::Pending;
                }
            }
        }
    }
}

impl<S, K: Clone> FairQueue<S, K> {
    pub fn new(block_on_no_clients: bool) -> Self {
        Self {
            block_on_no_clients,
            inner: Arc::new(Mutex::new(QueueInner {
                counter: atomic::AtomicUsize::new(0),
                ready_queue: BinaryHeap::new(),
                streams: HashMap::new(),
                waker: None,
            })),
        }
    }

    pub(crate) fn inner(&self) -> Arc<Mutex<QueueInner<S, K>>> {
        self.inner.clone()
    }
}

#[cfg(test)]
mod test {
    use crate::fair_queue::FairQueue;
    use futures::StreamExt;

    #[tokio::test]
    async fn test_fair_queue_ready() {
        let a = futures::stream::iter(vec!["a1", "a2", "a3"]);
        let b = futures::stream::iter(vec!["b1", "b2", "b3"]);
        let c = futures::stream::iter(vec!["c1", "c2", "c3"]);

        let mut f_queue: FairQueue<_, u64> = FairQueue::new(false);
        {
            let inner = f_queue.inner();
            let mut inner_lock = inner.lock();
            inner_lock.insert(1, a);
            inner_lock.insert(2, b);
            inner_lock.insert(3, c);
        }

        let mut results = Vec::new();
        while let Some(i) = f_queue.next().await {
            results.push(i);
        }
        assert_eq!(
            results,
            vec![
                (1, "a1"),
                (2, "b1"),
                (3, "c1"),
                (1, "a2"),
                (2, "b2"),
                (3, "c2"),
                (1, "a3"),
                (2, "b3"),
                (3, "c3")
            ]
        );
    }

    #[tokio::test]
    async fn test_fair_queue_different_size() {
        let a = futures::stream::iter(vec!["a1", "a2", "a3"]);
        let b = futures::stream::iter(vec!["b1"]);
        let c = futures::stream::iter(vec!["c1", "c2"]);

        let mut f_queue: FairQueue<_, u64> = FairQueue::new(false);
        {
            let inner = f_queue.inner();
            let mut inner_lock = inner.lock();
            inner_lock.insert(1, a);
            inner_lock.insert(2, b);
            inner_lock.insert(3, c);
        }

        let mut results = Vec::new();
        while let Some(i) = f_queue.next().await {
            results.push(i);
        }
        assert_eq!(
            results,
            vec![
                (1, "a1"),
                (2, "b1"),
                (3, "c1"),
                (1, "a2"),
                (3, "c2"),
                (1, "a3")
            ]
        );
    }
}
<|endoftext|>


<commit_before>use std::cell::UnsafeCell;
use std::ptr;
use std::sync::atomic::{AtomicUsize, fence, Ordering};

use super::MAX_JOBS;
use super::job::Job;

const MASK: usize = MAX_JOBS - 1;

/// A double-ended job queue.
/// This is currently locking, but will be made lock-free when performance demands it.
pub struct Queue {
    buf: UnsafeCell<[*mut Job; MAX_JOBS]>,
    top: AtomicUsize,
    bottom: AtomicUsize,
}

unsafe impl Send for Queue {}
unsafe impl Sync for Queue {}

macro_rules! memory_barrier {
    () => {
        fence(Ordering::SeqCst);
    }
}

#[cfg(nightly)]
macro_rules! compiler_barrier {
    () => {
        ::std::intrinsics::atomic_singlethreadfence()
    }
}

// do a full memory barrier on stable/beta.
#[cfg(not(nightly))]
macro_rules! compiler_barrier {
    () => {
        fence(Ordering::SeqCst);
    }
}

impl Queue {
    pub fn new() -> Self {
        Queue {
            buf: UnsafeCell::new([ptr::null_mut(); MAX_JOBS]),
            top: AtomicUsize::new(0),
            bottom: AtomicUsize::new(0),
        }
    }

    // push a job onto the private end of the queue.
    // this must be performed from the thread which logically
    // owns this queue.
    pub unsafe fn push(&self, job: *mut Job) {
        assert!(!job.is_null(), "Attempted to push null job onto queue");
        let b = self.bottom.load(Ordering::Acquire);
        (*self.buf.get())[b & MASK] = job;

        compiler_barrier!();
        self.bottom.store(b + 1, Ordering::Release);
    }

    // steal a job from the public end of the queue.
    pub fn steal(&self) -> Option<*mut Job> {
        unsafe {
            let t = self.top.load(Ordering::Acquire);

            compiler_barrier!();
            let b = self.bottom.load(Ordering::Acquire);
            if t < b {
                // non-empty queue.
                // may be racing against other steal/pop calls.
                let job = (*self.buf.get())[t & MASK];
                if self.top.compare_and_swap(t, t + 1, Ordering::AcqRel) != t {
                    // lost the race.
                    None
                } else {
                    // won the race.
                    Some(job)
                }
            } else {
                // empty queue.
                None
            }
        }
    }

    // pop a job from the private end of the queue.
    // this must be done from the thread which logically
    // owns this queue.
    pub unsafe fn pop(&self) -> Option<*mut Job> {
        let b = self.bottom.load(Ordering::Acquire).wrapping_sub(1);
        // check that bottom is not 0 before storing new b.
        if b == ::std::usize::MAX { return None }
        self.bottom.store(b, Ordering::Release);
        memory_barrier!();

        let t = self.top.load(Ordering::Acquire);
        if t <= b {
            // non-empty queue.
            let mut job = Some((*self.buf.get())[b & MASK]);
            if t != b {
                job
            } else {
                if self.top.compare_and_swap(t, t + 1, Ordering::SeqCst) != t {
                    job = None;
                }

                self.bottom.store(t + 1, Ordering::Release);
                job
            }
        } else {
            // empty queue.
            self.bottom.store(t, Ordering::Release);
            None
        }
    }

    // get the length of this queue.
    pub fn len(&self) -> usize {
       self.bottom.load(Ordering::Acquire) -  self.top.load(Ordering::Acquire)
    }

    // reset the counters, so they don't wrap around.
    pub fn reset_counters(&self) {
        debug_assert_eq!(self.len(), 0);

        self.bottom.store(0, Ordering::Release);
        self.top.store(0, Ordering::Release);
    }
}<commit_msg>steal now attempts in a loop until target is empty.<commit_after>use std::cell::UnsafeCell;
use std::ptr;
use std::sync::atomic::{AtomicUsize, fence, Ordering};

use super::MAX_JOBS;
use super::job::Job;

const MASK: usize = MAX_JOBS - 1;

/// A double-ended job queue.
/// This is currently locking, but will be made lock-free when performance demands it.
pub struct Queue {
    buf: UnsafeCell<[*mut Job; MAX_JOBS]>,
    top: AtomicUsize,
    bottom: AtomicUsize,
}

unsafe impl Send for Queue {}
unsafe impl Sync for Queue {}

macro_rules! memory_barrier {
    () => {
        fence(Ordering::SeqCst);
    }
}

#[cfg(nightly)]
macro_rules! compiler_barrier {
    () => {
        ::std::intrinsics::atomic_singlethreadfence()
    }
}

// do a full memory barrier on stable/beta.
#[cfg(not(nightly))]
macro_rules! compiler_barrier {
    () => {
        fence(Ordering::SeqCst);
    }
}

impl Queue {
    pub fn new() -> Self {
        Queue {
            buf: UnsafeCell::new([ptr::null_mut(); MAX_JOBS]),
            top: AtomicUsize::new(0),
            bottom: AtomicUsize::new(0),
        }
    }

    // push a job onto the private end of the queue.
    // this must be performed from the thread which logically
    // owns this queue.
    pub unsafe fn push(&self, job: *mut Job) {
        assert!(!job.is_null(), "Attempted to push null job onto queue");
        let b = self.bottom.load(Ordering::Acquire);
        (*self.buf.get())[b & MASK] = job;

        compiler_barrier!();
        self.bottom.store(b + 1, Ordering::Release);
    }

    // steal a job from the public end of the queue.
    // loops until it wins a race or the queue is empty.
    pub fn steal(&self) -> Option<*mut Job> {
        loop {
            unsafe {
                let t = self.top.load(Ordering::Acquire);

                compiler_barrier!();
                let b = self.bottom.load(Ordering::Acquire);
                if t < b {
                    // non-empty queue.
                    // may be racing against other steal/pop calls.
                    let job = (*self.buf.get())[t & MASK];
                    if self.top.compare_and_swap(t, t + 1, Ordering::AcqRel) != t {
                        // lost the race.
                        continue;
                    } else {
                        // won the race.
                        return Some(job);
                    }
                } else {
                    // empty queue.
                    return None;
                }
            }
        }
    }

    // pop a job from the private end of the queue.
    // this must be done from the thread which logically
    // owns this queue.
    pub unsafe fn pop(&self) -> Option<*mut Job> {
        let b = self.bottom.load(Ordering::Acquire).wrapping_sub(1);
        // check that bottom is not 0 before storing new b.
        if b == ::std::usize::MAX { return None }
        self.bottom.store(b, Ordering::Release);
        memory_barrier!();

        let t = self.top.load(Ordering::Acquire);
        if t <= b {
            // non-empty queue.
            let mut job = Some((*self.buf.get())[b & MASK]);
            if t != b {
                job
            } else {
                if self.top.compare_and_swap(t, t + 1, Ordering::SeqCst) != t {
                    job = None;
                }

                self.bottom.store(t + 1, Ordering::Release);
                job
            }
        } else {
            // empty queue.
            self.bottom.store(t, Ordering::Release);
            None
        }
    }

    // get the length of this queue.
    pub fn len(&self) -> usize {
       self.bottom.load(Ordering::Acquire) -  self.top.load(Ordering::Acquire)
    }

    // reset the counters, so they don't wrap around.
    pub fn reset_counters(&self) {
        debug_assert_eq!(self.len(), 0);

        self.bottom.store(0, Ordering::Release);
        self.top.store(0, Ordering::Release);
    }
}<|endoftext|>


<commit_before>//! implements the board representation

use std::vec;
use std::collections::VecMap;

use super::color::{Color, White, Black};
use super::piece::{Piece, King, WHITE_PIECES, BLACK_PIECES, ALL_PIECES};
use super::square::{File, Rank, Square};
use super::bitboard::BitBoard;

#[deriving(Clone)]
pub struct Board {
    data: VecMap<BitBoard>,
    white_data: BitBoard,
    black_data: BitBoard,
    empty_data: BitBoard,
}

impl Board {
    pub fn new() -> Board {
        let mut ans = VecMap::new();
        for x in ALL_PIECES.iter() {
            ans.insert(*x as uint, BitBoard::new());
        }
        Board {
            data: ans,
            white_data: BitBoard::new(),
            black_data: BitBoard::new(),
            empty_data: BitBoard::new_full(),
        }
    }
    fn piece_data(&self, p: Piece) -> &BitBoard {
        self.data.get(&(p as uint)).unwrap()
    }
    fn piece_data_mut(&mut self, p: Piece) -> &mut BitBoard {
        self.data.get_mut(&(p as uint)).unwrap()
    }
    fn color_data(&self, c: Color) -> &BitBoard {
        match c {
            White => &self.white_data,
            Black => &self.black_data,
        }
    }
    fn color_data_mut(&mut self, c: Color) -> &mut BitBoard {
        match c {
            White => &mut self.white_data,
            Black => &mut self.black_data,
        }
    }
    pub fn at(&self, s: Square) -> Option<Piece> {
        for x in ALL_PIECES.iter() {
            if self.is_piece_at(*x, s) {
                return Some(*x);
            }
        }
        None
    }
    pub fn is_piece_at(&self, p: Piece, s: Square) -> bool {
        self.piece_data(p).at(s)
    }
    pub fn is_empty_at(&self, s: Square) -> bool {
        self.empty_data.at(s)
    }
    pub fn is_color_at(&self, s: Square, c: Color) -> bool {
        self.color_data(c).at(s)
    }
    pub fn set_at_mut(&mut self, s: Square, p: Piece) {
        debug_assert!(self.is_empty_at(s), "set_at_mut(), s = {}", s);
        self.piece_data_mut(p).set_at_mut(s);
        self.color_data_mut(p.color()).set_at_mut(s);
        self.empty_data.remove_at_mut(s);
    }
    pub fn remove_at_mut(&mut self, s: Square, p: Piece) {
        debug_assert!(self.is_piece_at(p, s), "remove_at_mut(), s = {}", s);
        self.piece_data_mut(p).remove_at_mut(s);
        self.color_data_mut(p.color()).remove_at_mut(s);
        self.empty_data.set_at_mut(s);
    }
    pub fn king_square(&self, c: Color) -> Square {
        let curr_king = Piece::new(c, King);
        let temp = self.iter().find( |x| x.0 == curr_king ).unwrap();
        temp.1
    }
    fn piece_vec(&self) -> Vec<(Piece,Square)> {
        let mut ans: Vec<(Piece, Square)> = Vec::new();
        for f in range::<u8>(0,8) {
            for r in range::<u8>(0,8) {
                let s = Square::new(File(f),Rank(r));
                let x = self.at(s);
                if x.is_some() {
                    ans.push((x.unwrap(), s));
                }
            }
        }
        ans
    }
    pub fn iter(&self) -> Iter { Iter(self.piece_vec().into_iter()) }
}

pub struct Iter(vec::IntoIter<(Piece,Square)>);
impl Iterator<(Piece,Square)> for Iter {
    fn next(&mut self) -> Option<(Piece,Square)> { self.0.next() }
    fn size_hint(&self) -> (uint, Option<uint>) { self.0.size_hint() }
}
impl DoubleEndedIterator<(Piece,Square)> for Iter {
    fn next_back(&mut self) -> Option<(Piece,Square)> { self.0.next_back() }
}
impl ExactSizeIterator<(Piece,Square)> for Iter {}
<commit_msg>Optimize board representation.<commit_after>//! implements the board representation

use std::vec;
use std::collections::VecMap;

use super::color::{Color, White, Black};
use super::piece::{Piece, King, WHITE_PIECES, BLACK_PIECES, ALL_PIECES};
use super::square::Square;
use super::bitboard::BitBoard;

#[deriving(Clone)]
pub struct Board {
    data: VecMap<BitBoard>,
    white_data: BitBoard,
    black_data: BitBoard,
    empty_data: BitBoard,
}

impl Board {
    pub fn new() -> Board {
        let mut ans = VecMap::new();
        for x in ALL_PIECES.iter() {
            ans.insert(*x as uint, BitBoard::new());
        }
        Board {
            data: ans,
            white_data: BitBoard::new(),
            black_data: BitBoard::new(),
            empty_data: BitBoard::new_full(),
        }
    }
    fn piece_data(&self, p: Piece) -> &BitBoard {
        self.data.get(&(p as uint)).unwrap()
    }
    fn piece_data_mut(&mut self, p: Piece) -> &mut BitBoard {
        self.data.get_mut(&(p as uint)).unwrap()
    }
    fn color_data(&self, c: Color) -> &BitBoard {
        match c {
            White => &self.white_data,
            Black => &self.black_data,
        }
    }
    fn color_data_mut(&mut self, c: Color) -> &mut BitBoard {
        match c {
            White => &mut self.white_data,
            Black => &mut self.black_data,
        }
    }
    pub fn at(&self, s: Square) -> Option<Piece> {
        for x in ALL_PIECES.iter() {
            if self.is_piece_at(*x, s) {
                return Some(*x);
            }
        }
        None
    }
    pub fn is_piece_at(&self, p: Piece, s: Square) -> bool {
        self.piece_data(p).at(s)
    }
    pub fn is_empty_at(&self, s: Square) -> bool {
        self.empty_data.at(s)
    }
    pub fn is_color_at(&self, s: Square, c: Color) -> bool {
        self.color_data(c).at(s)
    }
    pub fn set_at_mut(&mut self, s: Square, p: Piece) {
        debug_assert!(self.is_empty_at(s), "set_at_mut(), s = {}", s);
        self.piece_data_mut(p).set_at_mut(s);
        self.color_data_mut(p.color()).set_at_mut(s);
        self.empty_data.remove_at_mut(s);
    }
    pub fn remove_at_mut(&mut self, s: Square, p: Piece) {
        debug_assert!(self.is_piece_at(p, s), "remove_at_mut(), s = {}", s);
        self.piece_data_mut(p).remove_at_mut(s);
        self.color_data_mut(p.color()).remove_at_mut(s);
        self.empty_data.set_at_mut(s);
    }
    pub fn king_square(&self, c: Color) -> Square {
        let curr_king = Piece::new(c, King);
        self.piece_data(curr_king).iter().next().unwrap()
    }
    fn piece_vec(&self) -> Vec<(Piece,Square)> {
        let mut ans: Vec<(Piece, Square)> = Vec::new();
        for p in ALL_PIECES.iter() {
            ans.extend(self.piece_data(*p).iter().map( |s: Square| (*p, s)  ) );
        }
        ans
    }
    pub fn iter(&self) -> Iter { Iter(self.piece_vec().into_iter()) }
}

pub struct Iter(vec::IntoIter<(Piece,Square)>);
impl Iterator<(Piece,Square)> for Iter {
    fn next(&mut self) -> Option<(Piece,Square)> { self.0.next() }
    fn size_hint(&self) -> (uint, Option<uint>) { self.0.size_hint() }
}
impl DoubleEndedIterator<(Piece,Square)> for Iter {
    fn next_back(&mut self) -> Option<(Piece,Square)> { self.0.next_back() }
}
impl ExactSizeIterator<(Piece,Square)> for Iter {}
<|endoftext|>


<commit_before>pub use engine::actor::Actor;
pub use engine::game::Game;
pub use engine::tiles::Tile;
pub use engine::map::Map;
pub use engine::world::World;

mod actor;
mod game;
mod tiles;
mod map;
mod world;
<commit_msg>Don't fully qualify Engine re-exports<commit_after>pub use self::actor::Actor;
pub use self::game::Game;
pub use self::tiles::Tile;
pub use self::map::Map;
pub use self::world::World;

mod actor;
mod game;
mod tiles;
mod map;
mod world;
<|endoftext|>


<commit_before>//! The main event loop which performs I/O on the pseudoterminal
use std::borrow::Cow;
use std::collections::VecDeque;
use std::io::{self, ErrorKind, Write};
use std::fs::File;
use std::os::unix::io::AsRawFd;
use std::sync::Arc;

use mio::{self, Events, PollOpt, Ready};
use mio::unix::EventedFd;

use ansi;
use display;
use event;
use term::Term;
use util::thread;
use sync::FairMutex;

/// Messages that may be sent to the `EventLoop`
#[derive(Debug)]
pub enum Msg {
    /// Data that should be written to the pty
    Input(Cow<'static, [u8]>),

    /// Indicates that the `EventLoop` should shut down, as Alacritty is shutting down
    Shutdown
}

/// The main event!.. loop.
///
/// Handles all the pty I/O and runs the pty parser which updates terminal
/// state.
pub struct EventLoop<Io> {
    poll: mio::Poll,
    pty: Io,
    rx: mio::channel::Receiver<Msg>,
    tx: mio::channel::Sender<Msg>,
    terminal: Arc<FairMutex<Term>>,
    display: display::Notifier,
    ref_test: bool,
}

/// Helper type which tracks how much of a buffer has been written.
struct Writing {
    source: Cow<'static, [u8]>,
    written: usize,
}

/// All of the mutable state needed to run the event loop
///
/// Contains list of items to write, current write state, etc. Anything that
/// would otherwise be mutated on the `EventLoop` goes here.
pub struct State {
    write_list: VecDeque<Cow<'static, [u8]>>,
    writing: Option<Writing>,
    parser: ansi::Processor,
}

pub struct Notifier(pub ::mio::channel::Sender<Msg>);

impl event::Notify for Notifier {
    fn notify<B>(&mut self, bytes: B)
        where B: Into<Cow<'static, [u8]>>
    {
        let bytes = bytes.into();
        // terminal hangs if we send 0 bytes through.
        if bytes.len() == 0 {
            return
        }
        match self.0.send(Msg::Input(bytes)) {
            Ok(_) => (),
            Err(_) => panic!("expected send event loop msg"),
        }
    }
}

impl Default for State {
    fn default() -> State {
        State {
            write_list: VecDeque::new(),
            parser: ansi::Processor::new(),
            writing: None,
        }
    }
}

impl State {
    #[inline]
    fn ensure_next(&mut self) {
        if self.writing.is_none() {
            self.goto_next();
        }
    }

    #[inline]
    fn goto_next(&mut self) {
        self.writing = self.write_list
            .pop_front()
            .map(Writing::new);
    }

    #[inline]
    fn take_current(&mut self) -> Option<Writing> {
        self.writing.take()
    }

    #[inline]
    fn needs_write(&self) -> bool {
        self.writing.is_some() || !self.write_list.is_empty()
    }

    #[inline]
    fn set_current(&mut self, new: Option<Writing>) {
        self.writing = new;
    }
}

impl Writing {
    #[inline]
    fn new(c: Cow<'static, [u8]>) -> Writing {
        Writing { source: c, written: 0 }
    }

    #[inline]
    fn advance(&mut self, n: usize) {
        self.written += n;
    }

    #[inline]
    fn remaining_bytes(&self) -> &[u8] {
        &self.source[self.written..]
    }

    #[inline]
    fn finished(&self) -> bool {
        self.written >= self.source.len()
    }
}

/// `mio::Token` for the event loop channel
const CHANNEL: mio::Token = mio::Token(0);

/// `mio::Token` for the pty file descriptor
const PTY: mio::Token = mio::Token(1);

impl<Io> EventLoop<Io>
    where Io: io::Read + io::Write + Send + AsRawFd + 'static
{
    /// Create a new event loop
    pub fn new(
        terminal: Arc<FairMutex<Term>>,
        display: display::Notifier,
        pty: Io,
        ref_test: bool,
    ) -> EventLoop<Io> {
        let (tx, rx) = ::mio::channel::channel();
        EventLoop {
            poll: mio::Poll::new().expect("create mio Poll"),
            pty: pty,
            tx: tx,
            rx: rx,
            terminal: terminal,
            display: display,
            ref_test: ref_test,
        }
    }

    pub fn channel(&self) -> mio::channel::Sender<Msg> {
        self.tx.clone()
    }

    // Drain the channel
    //
    // Returns `Ok` if the `EventLoop` should continue running.
    // `Ok(true)`is returned if items were received
    //
    // An `Err` indicates that the event loop should shut down
    fn drain_recv_channel(&self, state: &mut State) -> Result<bool, ()> {
        let mut received_item = false;
        while let Ok(msg) = self.rx.try_recv() {
            received_item = true;
            match msg {
                Msg::Input(input) => {
                    state.write_list.push_back(input);
                },
                Msg::Shutdown => {
                    return Err(())
                }
            }
        }

        Ok(received_item)
    }

    // Returns a `bool` indicating whether or not the event loop should continue running
    #[inline]
    fn channel_event(&mut self, state: &mut State) -> bool {
        if self.drain_recv_channel(state).is_err() {
            return false;
        }

        self.poll.reregister(
            &self.rx, CHANNEL,
            Ready::readable(),
            PollOpt::edge() | PollOpt::oneshot()
        ).expect("reregister channel");

        if state.needs_write() {
            self.poll.reregister(
                &EventedFd(&self.pty.as_raw_fd()),
                PTY,
                Ready::readable() | Ready::writable(),
                PollOpt::edge() | PollOpt::oneshot()
            ).expect("reregister fd after channel recv");
        }
        true
    }

    #[inline]
    fn pty_read<W>(
        &mut self,
        state: &mut State,
        buf: &mut [u8],
        mut writer: Option<&mut W>
    ) -> io::Result<()>
        where W: Write
    {
        loop {
            match self.pty.read(&mut buf[..]) {
                Ok(0) => break,
                Ok(got) => {
                    writer = writer.map(|w| {
                        w.write_all(&buf[..got]).unwrap(); w
                    });

                    let mut terminal = self.terminal.lock();
                    for byte in &buf[..got] {
                        state.parser.advance(&mut *terminal, *byte, &mut self.pty);
                    }

                    // Only request a draw if one hasn't already been requested.
                    //
                    // This is a performance optimization even if only for X11
                    // which is very expensive to hammer on the even loop wakeup
                    if !terminal.dirty {
                        self.display.notify();
                        terminal.dirty = true;

                        // Break for writing
                        //
                        // Want to prevent case where reading always returns
                        // data and sequences like `C-c` cannot be sent.
                        //
                        // Doing this check in !terminal.dirty will prevent the
                        // condition from being checked overzealously.
                        //
                        // We want to break if `drain_recv_channel` returns either `Ok(true)`
                        // (new items came in for writing) or `Err` (we need to shut down)
                        if state.writing.is_some()
                            || !state.write_list.is_empty()
                            || self.drain_recv_channel(state).unwrap_or_else(|_| true)
                        {
                            break;
                        }
                    }
                },
                Err(err) => {
                    match err.kind() {
                        ErrorKind::Interrupted |
                        ErrorKind::WouldBlock => break,
                        _ => return Err(err),
                    }
                }
            }
        }

        Ok(())
    }

    #[inline]
    fn pty_write(&mut self, state: &mut State) -> io::Result<()> {
        state.ensure_next();

        'write_many: while let Some(mut current) = state.take_current() {
            'write_one: loop {
                match self.pty.write(current.remaining_bytes()) {
                    Ok(0) => {
                        state.set_current(Some(current));
                        break 'write_many;
                    },
                    Ok(n) => {
                        current.advance(n);
                        if current.finished() {
                            state.goto_next();
                            break 'write_one;
                        }
                    },
                    Err(err) => {
                        state.set_current(Some(current));
                        match err.kind() {
                            ErrorKind::Interrupted |
                            ErrorKind::WouldBlock => break 'write_many,
                            _ => return Err(err),
                        }
                    }
                }

            }
        }

        Ok(())
    }

    pub fn spawn(
        mut self,
        state: Option<State>
    ) -> thread::JoinHandle<(EventLoop<Io>, State)> {
        thread::spawn_named("pty reader", move || {
            let mut state = state.unwrap_or_else(Default::default);
            let mut buf = [0u8; 4096];

            let fd = self.pty.as_raw_fd();
            let fd = EventedFd(&fd);

            let poll_opts = PollOpt::edge() | PollOpt::oneshot();

            self.poll.register(&self.rx, CHANNEL, Ready::readable(), poll_opts).unwrap();
            self.poll.register(&fd, PTY, Ready::readable(), poll_opts).unwrap();

            let mut events = Events::with_capacity(1024);

            let mut pipe = if self.ref_test {
                let file = File::create("./alacritty.recording")
                    .expect("create alacritty recording");
                Some(file)
            } else {
                None
            };

            'event_loop: loop {
                if let Err(err) = self.poll.poll(&mut events, None) {
                    match err.kind() {
                        ErrorKind::Interrupted => continue,
                        _ => panic!("EventLoop polling error: {:?}", err)
                    }
                }

                for event in events.iter() {
                    match event.token() {
                        CHANNEL =>  {
                            if !self.channel_event(&mut state) {
                                break 'event_loop;
                            }
                        },
                        PTY => {
                            let kind = event.kind();

                            if kind.is_hup() {
                                break 'event_loop;
                            }

                            if kind.is_readable() {
                                if let Err(err) = self.pty_read(&mut state, &mut buf, pipe.as_mut()) {
                                    error!("Event loop exitting due to error: {} [{}:{}]",
                                           err, file!(), line!());
                                    break 'event_loop;
                                }

                                if ::tty::process_should_exit() {
                                    break 'event_loop;
                                }
                            }

                            if kind.is_writable() {
                                if let Err(err) = self.pty_write(&mut state) {
                                    error!("Event loop exitting due to error: {} [{}:{}]",
                                           err, file!(), line!());
                                    break 'event_loop;
                                }
                            }

                            // Figure out pty interest
                            let mut interest = Ready::readable();
                            if state.needs_write() {
                                interest.insert(Ready::writable());
                            }

                            // Reregister pty
                            self.poll
                                .reregister(&fd, PTY, interest, poll_opts)
                                .expect("register fd after read/write");
                        },
                        _ => (),
                    }
                }
            }

            let _ = self.poll.deregister(&self.rx);
            let _ = self.poll.deregister(&fd);

            (self, state)
        })
    }
}
<commit_msg>Add DrainResult enum<commit_after>//! The main event loop which performs I/O on the pseudoterminal
use std::borrow::Cow;
use std::collections::VecDeque;
use std::io::{self, ErrorKind, Write};
use std::fs::File;
use std::os::unix::io::AsRawFd;
use std::sync::Arc;

use mio::{self, Events, PollOpt, Ready};
use mio::unix::EventedFd;

use ansi;
use display;
use event;
use term::Term;
use util::thread;
use sync::FairMutex;

/// Messages that may be sent to the `EventLoop`
#[derive(Debug)]
pub enum Msg {
    /// Data that should be written to the pty
    Input(Cow<'static, [u8]>),

    /// Indicates that the `EventLoop` should shut down, as Alacritty is shutting down
    Shutdown
}

/// The main event!.. loop.
///
/// Handles all the pty I/O and runs the pty parser which updates terminal
/// state.
pub struct EventLoop<Io> {
    poll: mio::Poll,
    pty: Io,
    rx: mio::channel::Receiver<Msg>,
    tx: mio::channel::Sender<Msg>,
    terminal: Arc<FairMutex<Term>>,
    display: display::Notifier,
    ref_test: bool,
}

/// Helper type which tracks how much of a buffer has been written.
struct Writing {
    source: Cow<'static, [u8]>,
    written: usize,
}

/// Indicates the result of draining the mio channel
#[derive(Debug)]
enum DrainResult {
    /// At least one new item was received
    ReceivedItem,
    /// Nothing was available to receive
    Empty,
    /// A shutdown message was received
    Shutdown
}

impl DrainResult {

    pub fn is_shutdown(&self) -> bool {
        match *self {
            DrainResult::Shutdown => true,
            _ => false
        }
    }

    #[allow(dead_code)]
    pub fn is_empty(&self) -> bool {
        match *self {
            DrainResult::Empty => true,
            _ => false
        }
    }

    #[allow(dead_code)]
    pub fn is_item(&self) -> bool {
        match *self {
            DrainResult::ReceivedItem => true,
            _ => false
        }
    }

}

/// All of the mutable state needed to run the event loop
///
/// Contains list of items to write, current write state, etc. Anything that
/// would otherwise be mutated on the `EventLoop` goes here.
pub struct State {
    write_list: VecDeque<Cow<'static, [u8]>>,
    writing: Option<Writing>,
    parser: ansi::Processor,
}

pub struct Notifier(pub ::mio::channel::Sender<Msg>);

impl event::Notify for Notifier {
    fn notify<B>(&mut self, bytes: B)
        where B: Into<Cow<'static, [u8]>>
    {
        let bytes = bytes.into();
        // terminal hangs if we send 0 bytes through.
        if bytes.len() == 0 {
            return
        }
        match self.0.send(Msg::Input(bytes)) {
            Ok(_) => (),
            Err(_) => panic!("expected send event loop msg"),
        }
    }
}

impl Default for State {
    fn default() -> State {
        State {
            write_list: VecDeque::new(),
            parser: ansi::Processor::new(),
            writing: None,
        }
    }
}

impl State {
    #[inline]
    fn ensure_next(&mut self) {
        if self.writing.is_none() {
            self.goto_next();
        }
    }

    #[inline]
    fn goto_next(&mut self) {
        self.writing = self.write_list
            .pop_front()
            .map(Writing::new);
    }

    #[inline]
    fn take_current(&mut self) -> Option<Writing> {
        self.writing.take()
    }

    #[inline]
    fn needs_write(&self) -> bool {
        self.writing.is_some() || !self.write_list.is_empty()
    }

    #[inline]
    fn set_current(&mut self, new: Option<Writing>) {
        self.writing = new;
    }
}

impl Writing {
    #[inline]
    fn new(c: Cow<'static, [u8]>) -> Writing {
        Writing { source: c, written: 0 }
    }

    #[inline]
    fn advance(&mut self, n: usize) {
        self.written += n;
    }

    #[inline]
    fn remaining_bytes(&self) -> &[u8] {
        &self.source[self.written..]
    }

    #[inline]
    fn finished(&self) -> bool {
        self.written >= self.source.len()
    }
}

/// `mio::Token` for the event loop channel
const CHANNEL: mio::Token = mio::Token(0);

/// `mio::Token` for the pty file descriptor
const PTY: mio::Token = mio::Token(1);

impl<Io> EventLoop<Io>
    where Io: io::Read + io::Write + Send + AsRawFd + 'static
{
    /// Create a new event loop
    pub fn new(
        terminal: Arc<FairMutex<Term>>,
        display: display::Notifier,
        pty: Io,
        ref_test: bool,
    ) -> EventLoop<Io> {
        let (tx, rx) = ::mio::channel::channel();
        EventLoop {
            poll: mio::Poll::new().expect("create mio Poll"),
            pty: pty,
            tx: tx,
            rx: rx,
            terminal: terminal,
            display: display,
            ref_test: ref_test,
        }
    }

    pub fn channel(&self) -> mio::channel::Sender<Msg> {
        self.tx.clone()
    }

    // Drain the channel
    //
    // Returns a `DrainResult` indicating the result of receiving from the channe;
    //
    fn drain_recv_channel(&self, state: &mut State) -> DrainResult {
        let mut received_item = false;
        while let Ok(msg) = self.rx.try_recv() {
            received_item = true;
            match msg {
                Msg::Input(input) => {
                    state.write_list.push_back(input);
                },
                Msg::Shutdown => {
                    return DrainResult::Shutdown;
                }
            }
        }

        return if received_item { DrainResult::ReceivedItem } else { DrainResult::Empty };
    }

    // Returns a `bool` indicating whether or not the event loop should continue running
    #[inline]
    fn channel_event(&mut self, state: &mut State) -> bool {
        if self.drain_recv_channel(state).is_shutdown() {
            return false;
        }

        self.poll.reregister(
            &self.rx, CHANNEL,
            Ready::readable(),
            PollOpt::edge() | PollOpt::oneshot()
        ).expect("reregister channel");

        if state.needs_write() {
            self.poll.reregister(
                &EventedFd(&self.pty.as_raw_fd()),
                PTY,
                Ready::readable() | Ready::writable(),
                PollOpt::edge() | PollOpt::oneshot()
            ).expect("reregister fd after channel recv");
        }
        true
    }

    #[inline]
    fn pty_read<W>(
        &mut self,
        state: &mut State,
        buf: &mut [u8],
        mut writer: Option<&mut W>
    ) -> io::Result<()>
        where W: Write
    {
        loop {
            match self.pty.read(&mut buf[..]) {
                Ok(0) => break,
                Ok(got) => {
                    writer = writer.map(|w| {
                        w.write_all(&buf[..got]).unwrap(); w
                    });

                    let mut terminal = self.terminal.lock();
                    for byte in &buf[..got] {
                        state.parser.advance(&mut *terminal, *byte, &mut self.pty);
                    }

                    // Only request a draw if one hasn't already been requested.
                    //
                    // This is a performance optimization even if only for X11
                    // which is very expensive to hammer on the even loop wakeup
                    if !terminal.dirty {
                        self.display.notify();
                        terminal.dirty = true;

                        // Break for writing
                        //
                        // Want to prevent case where reading always returns
                        // data and sequences like `C-c` cannot be sent.
                        //
                        // Doing this check in !terminal.dirty will prevent the
                        // condition from being checked overzealously.
                        //
                        // We want to break if `drain_recv_channel` returns either `Ok(true)`
                        // (new items came in for writing) or `Err` (we need to shut down)
                        if state.writing.is_some()
                            || !state.write_list.is_empty()
                            || match self.drain_recv_channel(state) {
                                DrainResult::Shutdown | DrainResult::ReceivedItem => true,
                                _ => false
                            }
                        {
                            break;
                        }
                    }
                },
                Err(err) => {
                    match err.kind() {
                        ErrorKind::Interrupted |
                        ErrorKind::WouldBlock => break,
                        _ => return Err(err),
                    }
                }
            }
        }

        Ok(())
    }

    #[inline]
    fn pty_write(&mut self, state: &mut State) -> io::Result<()> {
        state.ensure_next();

        'write_many: while let Some(mut current) = state.take_current() {
            'write_one: loop {
                match self.pty.write(current.remaining_bytes()) {
                    Ok(0) => {
                        state.set_current(Some(current));
                        break 'write_many;
                    },
                    Ok(n) => {
                        current.advance(n);
                        if current.finished() {
                            state.goto_next();
                            break 'write_one;
                        }
                    },
                    Err(err) => {
                        state.set_current(Some(current));
                        match err.kind() {
                            ErrorKind::Interrupted |
                            ErrorKind::WouldBlock => break 'write_many,
                            _ => return Err(err),
                        }
                    }
                }

            }
        }

        Ok(())
    }

    pub fn spawn(
        mut self,
        state: Option<State>
    ) -> thread::JoinHandle<(EventLoop<Io>, State)> {
        thread::spawn_named("pty reader", move || {
            let mut state = state.unwrap_or_else(Default::default);
            let mut buf = [0u8; 4096];

            let fd = self.pty.as_raw_fd();
            let fd = EventedFd(&fd);

            let poll_opts = PollOpt::edge() | PollOpt::oneshot();

            self.poll.register(&self.rx, CHANNEL, Ready::readable(), poll_opts).unwrap();
            self.poll.register(&fd, PTY, Ready::readable(), poll_opts).unwrap();

            let mut events = Events::with_capacity(1024);

            let mut pipe = if self.ref_test {
                let file = File::create("./alacritty.recording")
                    .expect("create alacritty recording");
                Some(file)
            } else {
                None
            };

            'event_loop: loop {
                if let Err(err) = self.poll.poll(&mut events, None) {
                    match err.kind() {
                        ErrorKind::Interrupted => continue,
                        _ => panic!("EventLoop polling error: {:?}", err)
                    }
                }

                for event in events.iter() {
                    match event.token() {
                        CHANNEL =>  {
                            if !self.channel_event(&mut state) {
                                break 'event_loop;
                            }
                        },
                        PTY => {
                            let kind = event.kind();

                            if kind.is_hup() {
                                break 'event_loop;
                            }

                            if kind.is_readable() {
                                if let Err(err) = self.pty_read(&mut state, &mut buf, pipe.as_mut()) {
                                    error!("Event loop exitting due to error: {} [{}:{}]",
                                           err, file!(), line!());
                                    break 'event_loop;
                                }

                                if ::tty::process_should_exit() {
                                    break 'event_loop;
                                }
                            }

                            if kind.is_writable() {
                                if let Err(err) = self.pty_write(&mut state) {
                                    error!("Event loop exitting due to error: {} [{}:{}]",
                                           err, file!(), line!());
                                    break 'event_loop;
                                }
                            }

                            // Figure out pty interest
                            let mut interest = Ready::readable();
                            if state.needs_write() {
                                interest.insert(Ready::writable());
                            }

                            // Reregister pty
                            self.poll
                                .reregister(&fd, PTY, interest, poll_opts)
                                .expect("register fd after read/write");
                        },
                        _ => (),
                    }
                }
            }

            let _ = self.poll.deregister(&self.rx);
            let _ = self.poll.deregister(&fd);

            (self, state)
        })
    }
}
<|endoftext|>