<commit_before>use std::any::Any;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use rustc::middle::cstore::EncodedMetadata;
use rustc::mir::mono::{Linkage as RLinkage, Visibility};
use rustc::session::config::{DebugInfo, OutputType};
use rustc_codegen_ssa::back::linker::LinkerInfo;
use rustc_codegen_ssa::CrateInfo;

use crate::prelude::*;

use crate::backend::{Emit, WriteDebugInfo};

pub fn codegen_crate(
    tcx: TyCtxt<'_>,
    metadata: EncodedMetadata,
    need_metadata_module: bool,
) -> Box<dyn Any> {
    tcx.sess.abort_if_errors();

    if std::env::var("SHOULD_RUN").is_ok()
        && tcx.sess.crate_types.get().contains(&CrateType::Executable)
    {
        #[cfg(not(target_arch = "wasm32"))]
        let _: ! = run_jit(tcx);

        #[cfg(target_arch = "wasm32")]
        panic!("jit not supported on wasm");
    }

    run_aot(tcx, metadata, need_metadata_module)
}

#[cfg(not(target_arch = "wasm32"))]
fn run_jit(tcx: TyCtxt<'_>) -> ! {
    use cranelift_simplejit::{SimpleJITBackend, SimpleJITBuilder};

    // Rustc opens us without the RTLD_GLOBAL flag, so __cg_clif_global_atomic_mutex will not be
    // exported. We fix this by opening ourself again as global.
    // FIXME remove once atomic_shim is gone
    let cg_dylib = std::ffi::OsString::from(&tcx.sess.opts.debugging_opts.codegen_backend.as_ref().unwrap());
    std::mem::forget(libloading::os::unix::Library::open(Some(cg_dylib), libc::RTLD_NOW | libc::RTLD_GLOBAL).unwrap());

    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let mut jit_builder = SimpleJITBuilder::with_isa(
        crate::build_isa(tcx.sess, false),
        cranelift_module::default_libcall_names(),
    );
    jit_builder.symbols(imported_symbols);
    let mut jit_module: Module<SimpleJITBackend> = Module::new(jit_builder);
    assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());

    let 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 main_func_id = jit_module
        .declare_function("main", Linkage::Import, &sig)
        .unwrap();

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

    codegen_mono_items(tcx, &mut jit_module, None, mono_items);
    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module);
    crate::allocator::codegen(tcx, &mut jit_module);

    jit_module.finalize_definitions();

    tcx.sess.abort_if_errors();

    let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);

    println!("Rustc codegen cranelift will JIT run the executable, because the SHOULD_RUN env var is set");

    let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
        unsafe { ::std::mem::transmute(finalized_main) };

    let args = ::std::env::var("JIT_ARGS").unwrap_or_else(|_| String::new());
    let args = args
        .split(" ")
        .chain(Some(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string()))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
    // TODO: Rust doesn't care, but POSIX argv has a NULL sentinel at the end

    let ret = f(args.len() as c_int, argv.as_ptr());

    jit_module.finish();
    std::process::exit(ret);
}

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

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(LOCAL_CRATE);
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == 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;
        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(|(_idx, symbol)| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                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 run_aot(
    tcx: TyCtxt<'_>,
    metadata: EncodedMetadata,
    need_metadata_module: bool,
) -> Box<CodegenResults> {
    let new_module = |name: String| {
        let module = crate::backend::make_module(tcx.sess, name);
        assert_eq!(pointer_ty(tcx), module.target_config().pointer_type());
        module
    };

    fn emit_module<B: Backend>(
        tcx: TyCtxt<'_>,
        name: String,
        kind: ModuleKind,
        mut module: Module<B>,
        debug: Option<DebugContext>,
    ) -> CompiledModule
        where B::Product: Emit + WriteDebugInfo,
    {
            module.finalize_definitions();
            let mut product = module.finish();

            if let Some(mut debug) = debug {
                debug.emit(&mut product);
            }

            let tmp_file = tcx
                .output_filenames(LOCAL_CRATE)
                .temp_path(OutputType::Object, Some(&name));
            let obj = product.emit();
            std::fs::write(&tmp_file, obj).unwrap();
            CompiledModule {
                name,
                kind,
                object: Some(tmp_file),
                bytecode: None,
                bytecode_compressed: None,
            }
        };

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

    let mut module = new_module("some_file".to_string());

    let mut debug = if tcx.sess.opts.debuginfo != DebugInfo::None {
        let debug = DebugContext::new(
            tcx,
            module.target_config().pointer_type().bytes() as u8,
        );
        Some(debug)
    } else {
        None
    };

    codegen_mono_items(tcx, &mut module, debug.as_mut(), mono_items);
    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module);

    tcx.sess.abort_if_errors();

    let mut allocator_module = new_module("allocator_shim".to_string());
    let created_alloc_shim = crate::allocator::codegen(tcx, &mut allocator_module);

    rustc_incremental::assert_dep_graph(tcx);
    rustc_incremental::save_dep_graph(tcx);
    rustc_incremental::finalize_session_directory(tcx.sess, tcx.crate_hash(LOCAL_CRATE));

    let metadata_module = if need_metadata_module {
        let _timer = tcx.prof.generic_activity("codegen crate metadata");
        let (metadata_cgu_name, tmp_file) = tcx.sess.time("write compressed metadata", || {
            use rustc::mir::mono::CodegenUnitNameBuilder;

            let cgu_name_builder = &mut CodegenUnitNameBuilder::new(tcx);
            let metadata_cgu_name = cgu_name_builder
                .build_cgu_name(LOCAL_CRATE, &["crate"], Some("metadata"))
                .as_str()
                .to_string();

            let tmp_file = tcx
                .output_filenames(LOCAL_CRATE)
                .temp_path(OutputType::Metadata, Some(&metadata_cgu_name));

            let obj = crate::backend::with_object(tcx.sess, &metadata_cgu_name, |object| {
                crate::metadata::write_metadata(tcx, object);
            });

            std::fs::write(&tmp_file, obj).unwrap();

            (metadata_cgu_name, tmp_file)
        });

        Some(CompiledModule {
            name: metadata_cgu_name,
            kind: ModuleKind::Metadata,
            object: Some(tmp_file),
            bytecode: None,
            bytecode_compressed: None,
        })
    } else {
        None
    };

    Box::new(CodegenResults {
        crate_name: tcx.crate_name(LOCAL_CRATE),
        modules: vec![emit_module(
            tcx,
            "some_file".to_string(),
            ModuleKind::Regular,
            module,
            debug,
        )],
        allocator_module: if created_alloc_shim {
            Some(emit_module(
                tcx,
                "allocator_shim".to_string(),
                ModuleKind::Allocator,
                allocator_module,
                None,
            ))
        } else {
            None
        },
        metadata_module,
        crate_hash: tcx.crate_hash(LOCAL_CRATE),
        metadata,
        windows_subsystem: None, // Windows is not yet supported
        linker_info: LinkerInfo::new(tcx),
        crate_info: CrateInfo::new(tcx),
    })
}

fn codegen_mono_items<'tcx>(
    tcx: TyCtxt<'tcx>,
    module: &mut Module<impl Backend + 'static>,
    debug_context: Option<&mut DebugContext<'tcx>>,
    mono_items: FxHashMap<MonoItem<'tcx>, (RLinkage, Visibility)>,
) {
    let mut cx = CodegenCx::new(tcx, module, debug_context);

    time(tcx.sess, "codegen mono items", move || {
        tcx.sess.time("predefine functions", || {
            for (&mono_item, &(linkage, visibility)) in &mono_items {
                match mono_item {
                    MonoItem::Fn(instance) => {
                        let (name, sig) =
                            get_function_name_and_sig(tcx, cx.module.isa().triple(), instance, false);
                        let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
                        cx.module.declare_function(&name, linkage, &sig).unwrap();
                    }
                    MonoItem::Static(_) | MonoItem::GlobalAsm(_) => {}
                }
            }
        });

        for (mono_item, (linkage, visibility)) in mono_items {
            crate::unimpl::try_unimpl(tcx, || {
                let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
                trans_mono_item(&mut cx, mono_item, linkage);
            });
        }

        tcx.sess.time("finalize CodegenCx", || cx.finalize());
    });
}

fn trans_mono_item<'clif, 'tcx, B: Backend + 'static>(
    cx: &mut crate::CodegenCx<'clif, 'tcx, B>,
    mono_item: MonoItem<'tcx>,
    linkage: Linkage,
) {
    let tcx = cx.tcx;
    match mono_item {
        MonoItem::Fn(inst) => {
            let _inst_guard =
                PrintOnPanic(|| format!("{:?} {}", inst, tcx.symbol_name(inst).name.as_str()));
            debug_assert!(!inst.substs.needs_infer());
            let _mir_guard = PrintOnPanic(|| {
                match inst.def {
                    InstanceDef::Item(_)
                    | InstanceDef::DropGlue(_, _)
                    | InstanceDef::Virtual(_, _) => {
                        let mut mir = ::std::io::Cursor::new(Vec::new());
                        crate::rustc_mir::util::write_mir_pretty(
                            tcx,
                            Some(inst.def_id()),
                            &mut mir,
                        )
                        .unwrap();
                        String::from_utf8(mir.into_inner()).unwrap()
                    }
                    _ => {
                        // FIXME fix write_mir_pretty for these instances
                        format!("{:#?}", tcx.instance_mir(inst.def))
                    }
                }
            });

            cx.tcx.sess.time("codegen fn", || crate::base::trans_fn(cx, inst, linkage));
        }
        MonoItem::Static(def_id) => {
            crate::constant::codegen_static(&mut cx.constants_cx, def_id);
        }
        MonoItem::GlobalAsm(hir_id) => {
            let item = tcx.hir().expect_item(hir_id);
            if let rustc_hir::ItemKind::GlobalAsm(rustc_hir::GlobalAsm { asm }) = item.kind {
                // FIXME implement global asm using an external assembler
                if asm.as_str().contains("__rust_probestack") {
                    return;
                } else {
                    tcx
                        .sess
                        .fatal(&format!("Unimplemented global asm mono item \"{}\"", asm));
                }
            } else {
                bug!("Expected GlobalAsm found {:?}", item);
            }
        }
    }
}

fn time<R>(sess: &Session, name: &'static str, f: impl FnOnce() -> R) -> R {
    println!("[{}] start", name);
    let before = std::time::Instant::now();
    let res = sess.time(name, f);
    let after = std::time::Instant::now();
    println!("[{}] end time: {:?}", name, after - before);
    res
}
<commit_msg>Move time "codegen mono items" out of codegen_mono_items<commit_after>use std::any::Any;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use rustc::middle::cstore::EncodedMetadata;
use rustc::mir::mono::{Linkage as RLinkage, Visibility};
use rustc::session::config::{DebugInfo, OutputType};
use rustc_codegen_ssa::back::linker::LinkerInfo;
use rustc_codegen_ssa::CrateInfo;

use crate::prelude::*;

use crate::backend::{Emit, WriteDebugInfo};

pub fn codegen_crate(
    tcx: TyCtxt<'_>,
    metadata: EncodedMetadata,
    need_metadata_module: bool,
) -> Box<dyn Any> {
    tcx.sess.abort_if_errors();

    if std::env::var("SHOULD_RUN").is_ok()
        && tcx.sess.crate_types.get().contains(&CrateType::Executable)
    {
        #[cfg(not(target_arch = "wasm32"))]
        let _: ! = run_jit(tcx);

        #[cfg(target_arch = "wasm32")]
        panic!("jit not supported on wasm");
    }

    run_aot(tcx, metadata, need_metadata_module)
}

#[cfg(not(target_arch = "wasm32"))]
fn run_jit(tcx: TyCtxt<'_>) -> ! {
    use cranelift_simplejit::{SimpleJITBackend, SimpleJITBuilder};

    // Rustc opens us without the RTLD_GLOBAL flag, so __cg_clif_global_atomic_mutex will not be
    // exported. We fix this by opening ourself again as global.
    // FIXME remove once atomic_shim is gone
    let cg_dylib = std::ffi::OsString::from(&tcx.sess.opts.debugging_opts.codegen_backend.as_ref().unwrap());
    std::mem::forget(libloading::os::unix::Library::open(Some(cg_dylib), libc::RTLD_NOW | libc::RTLD_GLOBAL).unwrap());

    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let mut jit_builder = SimpleJITBuilder::with_isa(
        crate::build_isa(tcx.sess, false),
        cranelift_module::default_libcall_names(),
    );
    jit_builder.symbols(imported_symbols);
    let mut jit_module: Module<SimpleJITBackend> = Module::new(jit_builder);
    assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());

    let 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 main_func_id = jit_module
        .declare_function("main", Linkage::Import, &sig)
        .unwrap();

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

    time(tcx.sess, "codegen mono items", || {
        codegen_mono_items(tcx, &mut jit_module, None, mono_items);
    });
    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module);
    crate::allocator::codegen(tcx, &mut jit_module);

    jit_module.finalize_definitions();

    tcx.sess.abort_if_errors();

    let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);

    println!("Rustc codegen cranelift will JIT run the executable, because the SHOULD_RUN env var is set");

    let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
        unsafe { ::std::mem::transmute(finalized_main) };

    let args = ::std::env::var("JIT_ARGS").unwrap_or_else(|_| String::new());
    let args = args
        .split(" ")
        .chain(Some(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string()))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
    // TODO: Rust doesn't care, but POSIX argv has a NULL sentinel at the end

    let ret = f(args.len() as c_int, argv.as_ptr());

    jit_module.finish();
    std::process::exit(ret);
}

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

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(LOCAL_CRATE);
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == 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;
        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(|(_idx, symbol)| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                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 run_aot(
    tcx: TyCtxt<'_>,
    metadata: EncodedMetadata,
    need_metadata_module: bool,
) -> Box<CodegenResults> {
    let new_module = |name: String| {
        let module = crate::backend::make_module(tcx.sess, name);
        assert_eq!(pointer_ty(tcx), module.target_config().pointer_type());
        module
    };

    fn emit_module<B: Backend>(
        tcx: TyCtxt<'_>,
        name: String,
        kind: ModuleKind,
        mut module: Module<B>,
        debug: Option<DebugContext>,
    ) -> CompiledModule
        where B::Product: Emit + WriteDebugInfo,
    {
            module.finalize_definitions();
            let mut product = module.finish();

            if let Some(mut debug) = debug {
                debug.emit(&mut product);
            }

            let tmp_file = tcx
                .output_filenames(LOCAL_CRATE)
                .temp_path(OutputType::Object, Some(&name));
            let obj = product.emit();
            std::fs::write(&tmp_file, obj).unwrap();
            CompiledModule {
                name,
                kind,
                object: Some(tmp_file),
                bytecode: None,
                bytecode_compressed: None,
            }
        };

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

    let mut module = new_module("some_file".to_string());

    let mut debug = if tcx.sess.opts.debuginfo != DebugInfo::None {
        let debug = DebugContext::new(
            tcx,
            module.target_config().pointer_type().bytes() as u8,
        );
        Some(debug)
    } else {
        None
    };

    time(tcx.sess, "codegen mono items", || {
        codegen_mono_items(tcx, &mut module, debug.as_mut(), mono_items);
    });
    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module);

    tcx.sess.abort_if_errors();

    let mut allocator_module = new_module("allocator_shim".to_string());
    let created_alloc_shim = crate::allocator::codegen(tcx, &mut allocator_module);

    rustc_incremental::assert_dep_graph(tcx);
    rustc_incremental::save_dep_graph(tcx);
    rustc_incremental::finalize_session_directory(tcx.sess, tcx.crate_hash(LOCAL_CRATE));

    let metadata_module = if need_metadata_module {
        let _timer = tcx.prof.generic_activity("codegen crate metadata");
        let (metadata_cgu_name, tmp_file) = tcx.sess.time("write compressed metadata", || {
            use rustc::mir::mono::CodegenUnitNameBuilder;

            let cgu_name_builder = &mut CodegenUnitNameBuilder::new(tcx);
            let metadata_cgu_name = cgu_name_builder
                .build_cgu_name(LOCAL_CRATE, &["crate"], Some("metadata"))
                .as_str()
                .to_string();

            let tmp_file = tcx
                .output_filenames(LOCAL_CRATE)
                .temp_path(OutputType::Metadata, Some(&metadata_cgu_name));

            let obj = crate::backend::with_object(tcx.sess, &metadata_cgu_name, |object| {
                crate::metadata::write_metadata(tcx, object);
            });

            std::fs::write(&tmp_file, obj).unwrap();

            (metadata_cgu_name, tmp_file)
        });

        Some(CompiledModule {
            name: metadata_cgu_name,
            kind: ModuleKind::Metadata,
            object: Some(tmp_file),
            bytecode: None,
            bytecode_compressed: None,
        })
    } else {
        None
    };

    Box::new(CodegenResults {
        crate_name: tcx.crate_name(LOCAL_CRATE),
        modules: vec![emit_module(
            tcx,
            "some_file".to_string(),
            ModuleKind::Regular,
            module,
            debug,
        )],
        allocator_module: if created_alloc_shim {
            Some(emit_module(
                tcx,
                "allocator_shim".to_string(),
                ModuleKind::Allocator,
                allocator_module,
                None,
            ))
        } else {
            None
        },
        metadata_module,
        crate_hash: tcx.crate_hash(LOCAL_CRATE),
        metadata,
        windows_subsystem: None, // Windows is not yet supported
        linker_info: LinkerInfo::new(tcx),
        crate_info: CrateInfo::new(tcx),
    })
}

fn codegen_mono_items<'tcx>(
    tcx: TyCtxt<'tcx>,
    module: &mut Module<impl Backend + 'static>,
    debug_context: Option<&mut DebugContext<'tcx>>,
    mono_items: FxHashMap<MonoItem<'tcx>, (RLinkage, Visibility)>,
) {
    let mut cx = CodegenCx::new(tcx, module, debug_context);

    tcx.sess.time("predefine functions", || {
        for (&mono_item, &(linkage, visibility)) in &mono_items {
            match mono_item {
                MonoItem::Fn(instance) => {
                    let (name, sig) =
                        get_function_name_and_sig(tcx, cx.module.isa().triple(), instance, false);
                    let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
                    cx.module.declare_function(&name, linkage, &sig).unwrap();
                }
                MonoItem::Static(_) | MonoItem::GlobalAsm(_) => {}
            }
        }
    });

    for (mono_item, (linkage, visibility)) in mono_items {
        crate::unimpl::try_unimpl(tcx, || {
            let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
            trans_mono_item(&mut cx, mono_item, linkage);
        });
    }

    tcx.sess.time("finalize CodegenCx", || cx.finalize());
}

fn trans_mono_item<'clif, 'tcx, B: Backend + 'static>(
    cx: &mut crate::CodegenCx<'clif, 'tcx, B>,
    mono_item: MonoItem<'tcx>,
    linkage: Linkage,
) {
    let tcx = cx.tcx;
    match mono_item {
        MonoItem::Fn(inst) => {
            let _inst_guard =
                PrintOnPanic(|| format!("{:?} {}", inst, tcx.symbol_name(inst).name.as_str()));
            debug_assert!(!inst.substs.needs_infer());
            let _mir_guard = PrintOnPanic(|| {
                match inst.def {
                    InstanceDef::Item(_)
                    | InstanceDef::DropGlue(_, _)
                    | InstanceDef::Virtual(_, _) => {
                        let mut mir = ::std::io::Cursor::new(Vec::new());
                        crate::rustc_mir::util::write_mir_pretty(
                            tcx,
                            Some(inst.def_id()),
                            &mut mir,
                        )
                        .unwrap();
                        String::from_utf8(mir.into_inner()).unwrap()
                    }
                    _ => {
                        // FIXME fix write_mir_pretty for these instances
                        format!("{:#?}", tcx.instance_mir(inst.def))
                    }
                }
            });

            cx.tcx.sess.time("codegen fn", || crate::base::trans_fn(cx, inst, linkage));
        }
        MonoItem::Static(def_id) => {
            crate::constant::codegen_static(&mut cx.constants_cx, def_id);
        }
        MonoItem::GlobalAsm(hir_id) => {
            let item = tcx.hir().expect_item(hir_id);
            if let rustc_hir::ItemKind::GlobalAsm(rustc_hir::GlobalAsm { asm }) = item.kind {
                // FIXME implement global asm using an external assembler
                if asm.as_str().contains("__rust_probestack") {
                    return;
                } else {
                    tcx
                        .sess
                        .fatal(&format!("Unimplemented global asm mono item \"{}\"", asm));
                }
            } else {
                bug!("Expected GlobalAsm found {:?}", item);
            }
        }
    }
}

fn time<R>(sess: &Session, name: &'static str, f: impl FnOnce() -> R) -> R {
    println!("[{}] start", name);
    let before = std::time::Instant::now();
    let res = sess.time(name, f);
    let after = std::time::Instant::now();
    println!("[{}] end time: {:?}", name, after - before);
    res
}
<|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;
use std::thread::{self, JoinHandle};

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

const INFINITY: i16 = 0x7FFF;
const NEG_INFINITY: i16 = -0x7FFF;
const LOSE: i16 = -0x4000;
// Small contempt factor to discourage draws
const DRAW: i16 = 1;

pub enum AI {
    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() -> AI {
        AI::Idle
    }

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

    pub fn stop(&mut self) {
        if let AI::Thinking { stop_signal, .. } = self {
            stop_signal.store(true, Ordering::Relaxed);
            *self = AI::Idle;
        }
    }

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

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

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

        let prev_ai = mem::replace(self, AI::Idle);

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

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

            // 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();

            // 2-ply iterative deepening
            let mut moves: Vec<(Move, i16)> = board
                .generate_moves()
                .into_iter()
                .map(|mv| {
                    let mut new_board = board;
                    new_board.apply_move(&mv);

                    let score =
                        -alphabeta_negamax(&new_board, &mut board_list, NEG_INFINITY, INFINITY, 1);
                    (mv, score)
                })
                .collect();

            moves.sort_by(|&(_, a), &(_, b)| b.cmp(&a));

            let mut max_score = NEG_INFINITY;
            let mut best_move = None;
            for (mv, _) in moves {
                if stop_signal_clone.load(Ordering::Relaxed) {
                    return;
                }

                let mut new_board = board;
                new_board.apply_move(&mv);

                let score = -alphabeta_negamax(
                    &new_board,
                    &mut board_list,
                    NEG_INFINITY,
                    -max_score,
                    depth - 1,
                );
                if score > max_score {
                    max_score = score;
                    best_move = Some(mv);
                }
            }

            match best_move {
                Some(mv) => move_sender.send(mv).expect("AI failed to send Move"),
                None => panic!("AI failed to find a move"),
            }
        });

        *self = AI::Thinking {
            move_recv,
            stop_signal,
            handle,
        }
    }
}

fn alphabeta_negamax(
    board: &Board,
    // This list does not include the current board
    mut board_list: &mut Vec<Board>,
    mut alpha: i16,
    mut beta: i16,
    depth: u8,
) -> i16 {
    let set_ttable = |eval_type, score| {
        let zobrist = board.zobrist();
        TTable::set(
            zobrist,
            Entry {
                eval_type,
                depth,
                score,
                zobrist,
            },
        );
    };

    match board.outcome() {
        Outcome::Draw => {
            // This only works because the draw Outcome does not consider draw by repetition
            set_ttable(EvalType::Exact, DRAW);
            return DRAW;
        }
        Outcome::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(EvalType::Exact, score);
            return score;
        }
        Outcome::InProgress => {}
    }

    if let Some(entry) = TTable::get(board.zobrist()) {
        if board_list.len() >= 8 && board_list.iter().filter(|&&b| b == *board).count() >= 2 {
            return DRAW;
        }

        if entry.zobrist == board.zobrist() && entry.depth >= depth {
            match entry.eval_type {
                EvalType::Exact => {
                    return entry.score;
                }
                EvalType::Beta => {
                    if entry.score >= beta {
                        return entry.score;
                    }
                    beta = entry.score;
                }
            }
        }
    }

    if depth == 0 {
        evaluate(board)
    } else {
        let mut best_score = NEG_INFINITY;

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

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

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

            if score >= beta {
                set_ttable(EvalType::Beta, score);
                return beta;
            } else if score > alpha {
                alpha = score;
            }
        }
        set_ttable(EvalType::Exact, best_score);
        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>Make AI a struct which holds a Status enum<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;
use std::thread::{self, JoinHandle};

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

const INFINITY: i16 = 0x7FFF;
const NEG_INFINITY: i16 = -0x7FFF;
const LOSE: i16 = -0x4000;
// Small contempt factor to discourage draws
const DRAW: i16 = 1;

pub struct AI {
    status: Status,
}

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,
        }
    }

    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 let Some(_) = result {
            self.status = Status::Idle;
        }
        result
    }

    pub fn think(&mut self, board: Board, board_list: Vec<Board>, depth: u8) {
        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 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");
            }

            // 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();

            // 2-ply iterative deepening
            let mut moves: Vec<(Move, i16)> = board
                .generate_moves()
                .into_iter()
                .map(|mv| {
                    let mut new_board = board;
                    new_board.apply_move(&mv);

                    let score =
                        -alphabeta_negamax(&new_board, &mut board_list, NEG_INFINITY, INFINITY, 1);
                    (mv, score)
                })
                .collect();

            moves.sort_by(|&(_, a), &(_, b)| b.cmp(&a));

            let mut max_score = NEG_INFINITY;
            let mut best_move = None;
            for (mv, _) in moves {
                if stop_signal_clone.load(Ordering::Relaxed) {
                    return;
                }

                let mut new_board = board;
                new_board.apply_move(&mv);

                let score = -alphabeta_negamax(
                    &new_board,
                    &mut board_list,
                    NEG_INFINITY,
                    -max_score,
                    depth - 1,
                );
                if score > max_score {
                    max_score = score;
                    best_move = Some(mv);
                }
            }

            match best_move {
                Some(mv) => move_sender.send(mv).expect("AI failed to send Move"),
                None => panic!("AI failed to find a move"),
            }
        });

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

fn alphabeta_negamax(
    board: &Board,
    // This list does not include the current board
    mut board_list: &mut Vec<Board>,
    mut alpha: i16,
    mut beta: i16,
    depth: u8,
) -> i16 {
    let set_ttable = |eval_type, score| {
        let zobrist = board.zobrist();
        TTable::set(
            zobrist,
            Entry {
                eval_type,
                depth,
                score,
                zobrist,
            },
        );
    };

    match board.outcome() {
        Outcome::Draw => {
            // This only works because the draw Outcome does not consider draw by repetition
            set_ttable(EvalType::Exact, DRAW);
            return DRAW;
        }
        Outcome::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(EvalType::Exact, score);
            return score;
        }
        Outcome::InProgress => {}
    }

    if let Some(entry) = TTable::get(board.zobrist()) {
        if board_list.len() >= 8 && board_list.iter().filter(|&&b| b == *board).count() >= 2 {
            return DRAW;
        }

        if entry.zobrist == board.zobrist() && entry.depth >= depth {
            match entry.eval_type {
                EvalType::Exact => {
                    return entry.score;
                }
                EvalType::Beta => {
                    if entry.score >= beta {
                        return entry.score;
                    }
                    beta = entry.score;
                }
            }
        }
    }

    if depth == 0 {
        evaluate(board)
    } else {
        let mut best_score = NEG_INFINITY;

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

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

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

            if score >= beta {
                set_ttable(EvalType::Beta, score);
                return beta;
            } else if score > alpha {
                alpha = score;
            }
        }
        set_ttable(EvalType::Exact, best_score);
        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 std::collections::BTreeSet;
use std::fmt::{Display, Error, Formatter};
use action::{POSSIBLE_DIRECTIONS,POSSIBLE_MOVE_TYPES,Coordinate,Move};

pub const NUM_COLS: usize = 5;
pub const NUM_ROWS: usize = 4;

// Hacks for the three states of the lower-left of the board in JunSeok vs YeonSeung game
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum BoardType {
    Left,
    LeftOrMiddle,
    Middle,
    RightOrMiddle,
    Right,
}

impl BoardType {
    fn is_final(&self) -> bool {
        match *self {
            BoardType::Left => true,
            BoardType::Middle => true,
            BoardType::Right => true,
            BoardType::LeftOrMiddle => false,
            BoardType::RightOrMiddle => false,
        }
    }

    // Can a board of type `current` become this type?
    fn applies_to(&self, current: Option<BoardType>) -> bool {
        match (current, *self) {
            // a none board type can change to anything.
            (None, _) => true,
            // LeftOrMiddle can change to itself, left, or middle.
            (Some(BoardType::LeftOrMiddle), BoardType::LeftOrMiddle) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Left) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::LeftOrMiddle), _) => false,
            // RightOrMiddle can change to itself, right, or middle.
            (Some(BoardType::RightOrMiddle), BoardType::RightOrMiddle) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Right) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::RightOrMiddle), _) => false,
            // Left, Middle, Right can change to themselves only.
            (Some(BoardType::Left), BoardType::Left) => true,
            (Some(BoardType::Left), _) => false,
            (Some(BoardType::Middle), BoardType::Middle) => true,
            (Some(BoardType::Middle), _) => false,
            (Some(BoardType::Right), BoardType::Right) => true,
            (Some(BoardType::Right), _) => false,
        }
    }

    fn induced_by(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match *self {
            BoardType::Left          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            BoardType::LeftOrMiddle  => c == Coordinate{row: 1, col: 0},
            BoardType::Middle        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            BoardType::RightOrMiddle => c == Coordinate{row: 3, col: 1},
            BoardType::Right         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
        }
    }
}

#[derive(Copy, Clone)]
enum Orientation {
    UpDown,
    LeftRight,
    UpLeft,
    UpRight,
    DownLeft,
    DownRight,
}

impl Orientation {
    fn to_str(&self) -> &'static str {
        match *self {
            Orientation::UpDown => "║",
            Orientation::LeftRight => "═",
            Orientation::UpLeft => "╝",
            Orientation::UpRight => "╚",
            Orientation::DownLeft => "╗",
            Orientation::DownRight => "╔",
        }
    }
}

#[derive(Copy, Clone)]
enum OrientationOption {
    Fixed(Orientation),
    IfRight(Orientation, Orientation),
    IfLeft(Orientation, Orientation),
    LeftOrMiddle(Orientation, Orientation),
    RightOrMiddle(Orientation, Orientation),
}

impl OrientationOption {
    fn for_board(&self, board_type: Option<BoardType>) -> &'static str {
        let orientation = match (*self, board_type) {
            (OrientationOption::Fixed(orientation), _) => orientation,
            (OrientationOption::IfRight(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::IfRight(_, notright), _) => notright,
            (OrientationOption::IfLeft(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::IfLeft(_, notleft), _) => notleft,
            (OrientationOption::LeftOrMiddle(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::LeftOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::LeftOrMiddle(_, _), _) => unreachable!(),
            (OrientationOption::RightOrMiddle(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::RightOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::RightOrMiddle(_, _), _) => unreachable!(),
        };
        orientation.to_str()
    }
}

const ORIENTATIONS: [[OrientationOption; NUM_COLS]; NUM_ROWS] = [
    [
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::IfRight(Orientation::UpRight, Orientation::UpDown),
        OrientationOption::Fixed(Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
    [
        OrientationOption::LeftOrMiddle(Orientation::UpDown, Orientation::UpRight),
        OrientationOption::RightOrMiddle(Orientation::UpDown, Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::IfLeft(Orientation::LeftRight, Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
];

const POSSIBLE_BOARD_TYPES: [BoardType; 5] = [
    BoardType::Left,
    BoardType::LeftOrMiddle,
    BoardType::Middle,
    BoardType::RightOrMiddle,
    BoardType::Right,
];

pub type BoardArray = [[bool; NUM_COLS]; NUM_ROWS];
pub struct Board {
    board: BoardArray,
    board_type: Option<BoardType>,
}

impl Board {
    pub fn new(array: BoardArray, board_type: Option<BoardType>) -> Board {
        Board{board: array, board_type: board_type}
    }

    pub fn make_move(&mut self, m: Move) {
        if let Some(bt) = m.new_board_type {
            if !bt.applies_to(self.board_type) {
                panic!("Board type is {:?}, not compatible with {:?}", self.board_type, bt);
            }
            self.board_type = m.new_board_type
        }
        self.set_squares(m, true)
    }

    pub fn undo_move(&mut self, m: Move) {
        self.board_type = m.old_board_type;
        self.set_squares(m, false)
    }

    fn set_squares(&mut self, m: Move, mode: bool) {
        self.board[m.coord.row][m.coord.col] = mode;
        for other_space in m.extensions().iter() {
            self.board[other_space.row][other_space.col] = mode;
        }
    }

    fn occupied(&self, c: Coordinate) -> bool {
        self.board[c.row][c.col]
    }

    // This assesses whether a coordinate can be placed on the board,
    // given the current type of the board.
    fn compatible(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match self.board_type {
            Some(BoardType::Left)          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::LeftOrMiddle)  => c != Coordinate{row: 1, col: 1},
            Some(BoardType::Middle)        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::RightOrMiddle) => c != Coordinate{row: 3, col: 0},
            Some(BoardType::Right)         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
            None => true,
        }
    }

    fn frontier(&self) -> Vec<Coordinate> {
        let mut results = Vec::new();
        for row in 0..NUM_ROWS {
            for col in 0..NUM_COLS {
                let coord = Coordinate{row: row, col: col};
                if self.occupied(coord) || !self.compatible(coord) {
                    continue;
                }
                let have_neighbor = POSSIBLE_DIRECTIONS.iter().any(|dir| {
                    if let Some(dest) = coord.move_in(*dir, 1) { self.occupied(dest) } else { false }
                });
                if have_neighbor {
                    results.push(coord);
                }
            }
        }
        results
    }

    fn board_type_final(&self) -> bool {
        if let Some(x) = self.board_type { x.is_final() } else { false }
    }

    pub fn legal_moves(&self) -> Vec<Move> {
        let mut results = Vec::new();
        for frontier_space in self.frontier().iter() {
            for move_type in POSSIBLE_MOVE_TYPES.iter() {
                let mov = match Move::new(*frontier_space, *move_type, self.board_type) {
                    Some(x) => x,
                    None => continue,
                };
                let mut other_space_taken = false;
                let mut induces_board_type = frontier_space.induces_board_type();
                for other_space in mov.extensions().iter() {
                    if other_space.induces_board_type() {
                        induces_board_type = true;
                    }
                    if self.occupied(*other_space) || !self.compatible(*other_space) {
                        other_space_taken = true;
                        break;
                    }
                }
                if other_space_taken {
                    continue;
                }

                if !induces_board_type || self.board_type_final() {
                    results.push(mov);
                    continue;
                }

                let mut ok_board_types: BTreeSet<_> = POSSIBLE_BOARD_TYPES.iter().cloned().filter(|board_type| {
                    board_type.applies_to(self.board_type) &&
                        board_type.induced_by(*frontier_space) &&
                        mov.extensions().iter().all(|coord| board_type.induced_by(*coord))
                }).collect();

                // Dominated board types...
                if ok_board_types.contains(&BoardType::LeftOrMiddle) {
                    ok_board_types.remove(&BoardType::Left);
                    ok_board_types.remove(&BoardType::Middle);
                }
                if ok_board_types.contains(&BoardType::RightOrMiddle) {
                    ok_board_types.remove(&BoardType::Right);
                    ok_board_types.remove(&BoardType::Middle);
                }

                for board_type in ok_board_types.iter() {
                    results.push(mov.with_board_type(*board_type));
                }
            }
        }
        results
    }
}

impl Display for Board {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), Error> {
        // Print header row
        try!(formatter.write_str("    "));
        for i in 0..NUM_COLS {
            try!(write!(formatter, " {}", i));
        }
        try!(formatter.write_str("\n"));

        // Print top box border
        try!(formatter.write_str("    ┌─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┬─"));
        }
        try!(formatter.write_str("┐\n"));

        for (i, row) in self.board.iter().enumerate() {
            // Print cell content
            try!(write!(formatter, "{: >2}  │", i));
            for col in row.iter() {
                let chr = if *col { "x" } else { " " };
                try!(write!(formatter, "{}│", chr));
            }
            try!(formatter.write_str("\n"));

            // Print box border between rows
            if i != NUM_ROWS - 1 {
                try!(formatter.write_str("    ├─"));
                for _ in 0..NUM_COLS - 1 {
                    try!(formatter.write_str("┼─"));
                }
                try!(formatter.write_str("┤\n"));
            }
        }

        // Print bottom box border
        try!(formatter.write_str("    └─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┴─"));
        }
        try!(formatter.write_str("┘\n"));

        write!(formatter, "{:?}", self.board_type)
    }
}
<commit_msg>Board#fmt: Print orientations<commit_after>use std::collections::BTreeSet;
use std::fmt::{Display, Error, Formatter};
use action::{POSSIBLE_DIRECTIONS,POSSIBLE_MOVE_TYPES,Coordinate,Move};

pub const NUM_COLS: usize = 5;
pub const NUM_ROWS: usize = 4;

// Hacks for the three states of the lower-left of the board in JunSeok vs YeonSeung game
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub enum BoardType {
    Left,
    LeftOrMiddle,
    Middle,
    RightOrMiddle,
    Right,
}

impl BoardType {
    fn is_final(&self) -> bool {
        match *self {
            BoardType::Left => true,
            BoardType::Middle => true,
            BoardType::Right => true,
            BoardType::LeftOrMiddle => false,
            BoardType::RightOrMiddle => false,
        }
    }

    // Can a board of type `current` become this type?
    fn applies_to(&self, current: Option<BoardType>) -> bool {
        match (current, *self) {
            // a none board type can change to anything.
            (None, _) => true,
            // LeftOrMiddle can change to itself, left, or middle.
            (Some(BoardType::LeftOrMiddle), BoardType::LeftOrMiddle) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Left) => true,
            (Some(BoardType::LeftOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::LeftOrMiddle), _) => false,
            // RightOrMiddle can change to itself, right, or middle.
            (Some(BoardType::RightOrMiddle), BoardType::RightOrMiddle) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Right) => true,
            (Some(BoardType::RightOrMiddle), BoardType::Middle) => true,
            (Some(BoardType::RightOrMiddle), _) => false,
            // Left, Middle, Right can change to themselves only.
            (Some(BoardType::Left), BoardType::Left) => true,
            (Some(BoardType::Left), _) => false,
            (Some(BoardType::Middle), BoardType::Middle) => true,
            (Some(BoardType::Middle), _) => false,
            (Some(BoardType::Right), BoardType::Right) => true,
            (Some(BoardType::Right), _) => false,
        }
    }

    fn induced_by(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match *self {
            BoardType::Left          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            BoardType::LeftOrMiddle  => c == Coordinate{row: 1, col: 0},
            BoardType::Middle        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            BoardType::RightOrMiddle => c == Coordinate{row: 3, col: 1},
            BoardType::Right         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
        }
    }
}

#[derive(Copy, Clone)]
enum Orientation {
    UpDown,
    LeftRight,
    UpLeft,
    UpRight,
    DownLeft,
    DownRight,
}

impl Orientation {
    fn to_str(&self) -> &'static str {
        match *self {
            Orientation::UpDown => "║",
            Orientation::LeftRight => "═",
            Orientation::UpLeft => "╝",
            Orientation::UpRight => "╚",
            Orientation::DownLeft => "╗",
            Orientation::DownRight => "╔",
        }
    }
}

#[derive(Copy, Clone)]
enum OrientationOption {
    Fixed(Orientation),
    IfRight(Orientation, Orientation),
    IfLeft(Orientation, Orientation),
    LeftOrMiddle(Orientation, Orientation),
    RightOrMiddle(Orientation, Orientation),
}

impl OrientationOption {
    fn for_board(&self, board_type: Option<BoardType>) -> &'static str {
        let orientation = match (*self, board_type) {
            (OrientationOption::Fixed(orientation), _) => orientation,
            (OrientationOption::IfRight(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::IfRight(_, notright), _) => notright,
            (OrientationOption::IfLeft(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::IfLeft(_, notleft), _) => notleft,
            (OrientationOption::LeftOrMiddle(ifleft, _), Some(BoardType::Left)) => ifleft,
            (OrientationOption::LeftOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::LeftOrMiddle(_, _), _) => unreachable!(),
            (OrientationOption::RightOrMiddle(ifright, _), Some(BoardType::Right)) => ifright,
            (OrientationOption::RightOrMiddle(_, ifmiddle), Some(BoardType::Middle)) => ifmiddle,
            (OrientationOption::RightOrMiddle(_, _), _) => unreachable!(),
        };
        orientation.to_str()
    }
}

const ORIENTATIONS: [[OrientationOption; NUM_COLS]; NUM_ROWS] = [
    [
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::IfRight(Orientation::UpRight, Orientation::UpDown),
        OrientationOption::Fixed(Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::DownRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
    [
        OrientationOption::LeftOrMiddle(Orientation::UpDown, Orientation::UpRight),
        OrientationOption::RightOrMiddle(Orientation::UpDown, Orientation::DownLeft),
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::DownLeft)
    ],
    [
        OrientationOption::Fixed(Orientation::UpRight),
        OrientationOption::IfLeft(Orientation::LeftRight, Orientation::UpRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::LeftRight),
        OrientationOption::Fixed(Orientation::UpLeft)
    ],
];

const POSSIBLE_BOARD_TYPES: [BoardType; 5] = [
    BoardType::Left,
    BoardType::LeftOrMiddle,
    BoardType::Middle,
    BoardType::RightOrMiddle,
    BoardType::Right,
];

pub type BoardArray = [[bool; NUM_COLS]; NUM_ROWS];
pub struct Board {
    board: BoardArray,
    board_type: Option<BoardType>,
}

impl Board {
    pub fn new(array: BoardArray, board_type: Option<BoardType>) -> Board {
        Board{board: array, board_type: board_type}
    }

    pub fn make_move(&mut self, m: Move) {
        if let Some(bt) = m.new_board_type {
            if !bt.applies_to(self.board_type) {
                panic!("Board type is {:?}, not compatible with {:?}", self.board_type, bt);
            }
            self.board_type = m.new_board_type
        }
        self.set_squares(m, true)
    }

    pub fn undo_move(&mut self, m: Move) {
        self.board_type = m.old_board_type;
        self.set_squares(m, false)
    }

    fn set_squares(&mut self, m: Move, mode: bool) {
        self.board[m.coord.row][m.coord.col] = mode;
        for other_space in m.extensions().iter() {
            self.board[other_space.row][other_space.col] = mode;
        }
    }

    fn occupied(&self, c: Coordinate) -> bool {
        self.board[c.row][c.col]
    }

    // This assesses whether a coordinate can be placed on the board,
    // given the current type of the board.
    fn compatible(&self, c: Coordinate) -> bool {
        // Not in the lower left, so it's a free pass.
        if !c.induces_board_type() {
            return true;
        }

        match self.board_type {
            Some(BoardType::Left)          => c != Coordinate{row: 2, col: 1} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::LeftOrMiddle)  => c != Coordinate{row: 1, col: 1},
            Some(BoardType::Middle)        => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 1, col: 1},
            Some(BoardType::RightOrMiddle) => c != Coordinate{row: 3, col: 0},
            Some(BoardType::Right)         => c != Coordinate{row: 3, col: 0} && c != Coordinate{row: 2, col: 0},
            None => true,
        }
    }

    fn frontier(&self) -> Vec<Coordinate> {
        let mut results = Vec::new();
        for row in 0..NUM_ROWS {
            for col in 0..NUM_COLS {
                let coord = Coordinate{row: row, col: col};
                if self.occupied(coord) || !self.compatible(coord) {
                    continue;
                }
                let have_neighbor = POSSIBLE_DIRECTIONS.iter().any(|dir| {
                    if let Some(dest) = coord.move_in(*dir, 1) { self.occupied(dest) } else { false }
                });
                if have_neighbor {
                    results.push(coord);
                }
            }
        }
        results
    }

    fn board_type_final(&self) -> bool {
        if let Some(x) = self.board_type { x.is_final() } else { false }
    }

    pub fn legal_moves(&self) -> Vec<Move> {
        let mut results = Vec::new();
        for frontier_space in self.frontier().iter() {
            for move_type in POSSIBLE_MOVE_TYPES.iter() {
                let mov = match Move::new(*frontier_space, *move_type, self.board_type) {
                    Some(x) => x,
                    None => continue,
                };
                let mut other_space_taken = false;
                let mut induces_board_type = frontier_space.induces_board_type();
                for other_space in mov.extensions().iter() {
                    if other_space.induces_board_type() {
                        induces_board_type = true;
                    }
                    if self.occupied(*other_space) || !self.compatible(*other_space) {
                        other_space_taken = true;
                        break;
                    }
                }
                if other_space_taken {
                    continue;
                }

                if !induces_board_type || self.board_type_final() {
                    results.push(mov);
                    continue;
                }

                let mut ok_board_types: BTreeSet<_> = POSSIBLE_BOARD_TYPES.iter().cloned().filter(|board_type| {
                    board_type.applies_to(self.board_type) &&
                        board_type.induced_by(*frontier_space) &&
                        mov.extensions().iter().all(|coord| board_type.induced_by(*coord))
                }).collect();

                // Dominated board types...
                if ok_board_types.contains(&BoardType::LeftOrMiddle) {
                    ok_board_types.remove(&BoardType::Left);
                    ok_board_types.remove(&BoardType::Middle);
                }
                if ok_board_types.contains(&BoardType::RightOrMiddle) {
                    ok_board_types.remove(&BoardType::Right);
                    ok_board_types.remove(&BoardType::Middle);
                }

                for board_type in ok_board_types.iter() {
                    results.push(mov.with_board_type(*board_type));
                }
            }
        }
        results
    }
}

impl Display for Board {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), Error> {
        // Print header row
        try!(formatter.write_str("    "));
        for i in 0..NUM_COLS {
            try!(write!(formatter, " {}", i));
        }
        try!(formatter.write_str("\n"));

        // Print top box border
        try!(formatter.write_str("    ┌─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┬─"));
        }
        try!(formatter.write_str("┐\n"));

        for (i, row) in self.board.iter().enumerate() {
            // Print cell content
            try!(write!(formatter, "{: >2}  │", i));
            for (j, col) in row.iter().enumerate() {
                let chr = if *col { ORIENTATIONS[i][j].for_board(self.board_type) } else { " " };
                try!(write!(formatter, "{}│", chr));
            }
            try!(formatter.write_str("\n"));

            // Print box border between rows
            if i != NUM_ROWS - 1 {
                try!(formatter.write_str("    ├─"));
                for _ in 0..NUM_COLS - 1 {
                    try!(formatter.write_str("┼─"));
                }
                try!(formatter.write_str("┤\n"));
            }
        }

        // Print bottom box border
        try!(formatter.write_str("    └─"));
        for _ in 0..NUM_COLS - 1 {
            try!(formatter.write_str("┴─"));
        }
        try!(formatter.write_str("┘\n"));

        write!(formatter, "{:?}", self.board_type)
    }
}
<|endoftext|>


<commit_before>extern crate r2d2;

quick_error! {
    /// The most generic error message in ReQL
    #[derive(Debug)]
    pub enum Error {
        Compile(descr: String) { }
        Runtime(err: RuntimeError) { from() }
        Driver(err: DriverError) { from() }
    }
}

quick_error! {
    /// The parent class of all runtime errors
    ///
    /// All errors on the server unrelated to compilation. Programs may use this to catch any runtime
    /// error, but the server will always return a more specific error class.
    #[derive(Debug)]
    pub enum RuntimeError {
        QueryLogic(err: QueryLogicError) { from() }
        ResourceLimit(descr: String) {}
        User(descr: String) {}
        Internal(descr: String) {}
        Timeout(descr: String) {}
        Availability(err: AvailabilityError) { from() }
        Permissions(descr: String) {}
    }
}

quick_error! {
    /// The query contains a logical impossibility, such as adding a number to a string.
    #[derive(Debug)]
    pub enum QueryLogicError {
        NonExistence(descr: String) {}
    }
}

quick_error! {
    /// A server in the cluster is unavailable
    ///
    /// The parent class of `OpFailedError` and `OpIndeterminateError`. Programs may use this
    /// to catch any availability error, but the server will always return one of this class’s
    /// children.
    #[derive(Debug)]
    pub enum AvailabilityError {
        OpFailed(descr: String) {}
        OpIndeterminate(descr: String) {}
    }
}

quick_error! {
    /// An error has occurred within the driver
    ///
    /// This may be a driver bug, or it may be an unfulfillable command, such as an unserializable
    /// query.
    #[derive(Debug)]
    pub enum DriverError {
        Auth(descr: String) {}
        Initialization(err: r2d2::InitializationError) { from() }
        Connection(err: ConnectionError) { from() }
    }
}

quick_error! {
    /// Connection related errors
    #[derive(Debug)]
    pub enum ConnectionError {
        PoolWrite(descr: String) {}
        PoolRead(descr: String) {}
    }
}

/// Converts external error to ours
impl From<r2d2::InitializationError> for Error {
    fn from(err: r2d2::InitializationError) -> Error {
        From::from(DriverError::Initialization(err))
    }
}

impl From<ConnectionError> for Error {
    fn from(err: ConnectionError) -> Error {
        From::from(DriverError::Connection(err))
    }
}
<commit_msg>Convert IO error<commit_after>extern crate r2d2;

use std::io;

quick_error! {
    /// The most generic error message in ReQL
    #[derive(Debug)]
    pub enum Error {
        Compile(descr: String) { }
        Runtime(err: RuntimeError) { from() }
        Driver(err: DriverError) { from() }
    }
}

quick_error! {
    /// The parent class of all runtime errors
    ///
    /// All errors on the server unrelated to compilation. Programs may use this to catch any runtime
    /// error, but the server will always return a more specific error class.
    #[derive(Debug)]
    pub enum RuntimeError {
        QueryLogic(err: QueryLogicError) { from() }
        ResourceLimit(descr: String) {}
        User(descr: String) {}
        Internal(descr: String) {}
        Timeout(descr: String) {}
        Availability(err: AvailabilityError) { from() }
        Permissions(descr: String) {}
    }
}

quick_error! {
    /// The query contains a logical impossibility, such as adding a number to a string.
    #[derive(Debug)]
    pub enum QueryLogicError {
        NonExistence(descr: String) {}
    }
}

quick_error! {
    /// A server in the cluster is unavailable
    ///
    /// The parent class of `OpFailedError` and `OpIndeterminateError`. Programs may use this
    /// to catch any availability error, but the server will always return one of this class’s
    /// children.
    #[derive(Debug)]
    pub enum AvailabilityError {
        OpFailed(descr: String) {}
        OpIndeterminate(descr: String) {}
    }
}

quick_error! {
    /// An error has occurred within the driver
    ///
    /// This may be a driver bug, or it may be an unfulfillable command, such as an unserializable
    /// query.
    #[derive(Debug)]
    pub enum DriverError {
        Auth(descr: String) {}
        Initialization(err: r2d2::InitializationError) { from() }
        Connection(err: ConnectionError) { from() }
    }
}

quick_error! {
    /// Connection related errors
    #[derive(Debug)]
    pub enum ConnectionError {
        PoolWrite(descr: String) {}
        PoolRead(descr: String) {}
        Io(err: io::Error) { from() }
    }
}

/// Converts from r2d2 error
impl From<r2d2::InitializationError> for Error {
    fn from(err: r2d2::InitializationError) -> Error {
        From::from(DriverError::Initialization(err))
    }
}

impl From<ConnectionError> for Error {
    fn from(err: ConnectionError) -> Error {
        From::from(DriverError::Connection(err))
    }
}

/// Converts from IO error
impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        From::from(ConnectionError::Io(err))
    }
}
<|endoftext|>


<commit_before>use std::cast;
use std::mem::size_of;
use std::raw::Slice;
use std::num::Zero;

mod buffer;
mod schema;

static LEAF_MARKER: u8 = 0b111111111;
static BRANCH_MARKER: u8 = 0b0;

struct BTree<'a, K> {
	segment: u64,
	manager: &'a mut buffer::BufferManager,
	tree: LazyNode,
	next_free_page: u64,
}

impl<'a, K: TotalOrd + Zero> BTree<'a, K> {
	fn new<'b>(segment_id: u64, dummy: K, manager: &'b mut buffer::BufferManager) -> BTree<'b, K> {
		BTree {
			segment: segment_id,
			manager: manager,
			tree: LazyNode::new(1),
			next_free_page: 2,
		}
	}

	//fn locate_page

	fn insert(&mut self, key: K, value: schema::TID) {
		let node = self.tree.load(self.manager);
		// try insertion and see if the root was split
		let candidate = match node {
			Branch(mut n) => n.insert_value(self, key, value),
			Leaf(mut n) => n.insert_value(key, value),
		};
		// set new tree root if it was split
		match candidate {
			Some(new_node) => self.tree = new_node,
			None => (),
		}
	}

	fn next_page(&mut self) -> u64 {
		let n = self.next_free_page;
		self.next_free_page += 1;
		n
	}

	fn create_branch_page(&mut self) -> LazyNode {
		let next = self.next_page();
		let page_path = buffer::join_segment(self.segment, next);
		let pagelock = self.manager.fix_page(page_path).unwrap();
		let mut page = pagelock.write();
		let data = page.get_mut_data();
		data[0] = 0;
		LazyNode::new(page_path)
	}

	fn create_leaf_page(&mut self) -> LazyNode {
		let next = self.next_page();
		let page_path = buffer::join_segment(self.segment, next);
		let pagelock = self.manager.fix_page(page_path).unwrap();
		let mut page = pagelock.write();
		let data = page.get_mut_data();
		// marker to be a leaf page
		data[0] = 255;
		LazyNode::new(page_path)

	}

	fn erase(&mut self, key: K) {
	}

	fn lookup(self, key: K) -> Option<schema::TID> {
		None
	}
}

/* a placeholder for the actual page */
struct LazyNode {
	page_id: u64,
}

// TODO implement Drop to unfix page
impl LazyNode {
	fn new(page_id: u64) -> LazyNode {
		LazyNode {page_id: page_id}
	}

	/*
	 * As this is just a placeholder, return the actual node that his is
	 * representing
	 */
	fn load<'a, K: TotalOrd + Zero>(&self, manager: &mut buffer::BufferManager) -> Node<'a, K> {
		let pagelock = manager.fix_page(self.page_id).unwrap();
		let page = pagelock.read();
		LazyNode::load_node(page.get_data())
	}

	/* helper method to load a Node tepending on the page marker */
	fn load_node<'a, K: TotalOrd + Zero>(page: &[u8]) -> Node<'a, K> {
		if page[0] == 255 {
			Leaf(LeafNode::new(page))
		} else {
			Branch(BranchNode::new(page))
		}
	}
}

/* a node might either be an inner node (branch node) or a leaf node (LeafNode) */
enum Node<'a, K> {
	Branch(BranchNode<'a, K>),
	Leaf(LeafNode<'a, K>),
}

/* leaf pages consist mainly of leaf entries which are (K, TID) pairs */
struct LeafEntry<K> {
	key: K,
	tid: schema::TID,
}

/* branch pages consist mainly of branch entries which are (K, page_id) pairs */
struct BranchEntry<K> {
	key: K,
	page_id: u64,
}

struct LeafNode<'a, K> {
	capacity: uint,
	entries: &'a mut [LeafEntry<K>],
}

impl<'a, K: Zero> LeafNode<'a, K> {
	fn new(page: &[u8]) -> LeafNode<'a, K> {
		// first byte is a Leaf/Branch marker
		let entry_size = size_of::<LeafEntry<K>>();
		let entry_num = (page.len() - size_of::<u8>()) / entry_size;

		let mut r: &mut[LeafEntry<K>] = unsafe {
			cast::transmute(
				Slice::<LeafEntry<K>> {
					data: page.as_ptr() as *() as *LeafEntry<K>,
					len: entry_num,
				}
			)
		};

		/* find slots that are used */
		let mut capacity = r.len();
		for i in range(0, r.len()) {
			if !r[i].key.is_zero() && !r[i].tid.is_invalid() {
				capacity -= 1;
			}
		}

		LeafNode {
			entries: r,
			capacity: capacity
		}
	}

	fn insert_value(&mut self, key: K, tid: schema::TID) -> Option<LazyNode> {
		if self.capacity == 0 {
			// TODO split
			fail!("shit");
			return None
		}

		// find place to insert
		for i in range(0, self.entries.len()) {
			println!("Entry {:?}", self.entries[i]);
		}
		// TODO insert

		// insertion went fine, done
		None
	}
}

struct BranchNode<'a, K> {
	capacity: uint,
	entries: &'a mut [BranchEntry<K>],
}

impl<'a, K: TotalOrd + Zero> BranchNode<'a, K> {
	fn new(page: &[u8]) -> BranchNode<'a, K> {
		let entry_size = size_of::<BranchEntry<K>>();
		let entry_num = (page.len() - size_of::<u8>()) / entry_size;
		let mut r: &mut [BranchEntry<K>] = unsafe {
			cast::transmute(
				Slice::<BranchEntry<K>> {
					data: page.as_ptr() as *() as *BranchEntry<K>,
					len: entry_num,
				}
			)
		};

		/* find slots that are used */
		let mut capacity = r.len();
		for i in range(0, r.len()) {
			if r[i].page_id != 0 {
				capacity -= 1;
			}
		}

		println!("BranchNode computed capacity: {}", capacity);

		BranchNode {
			entries: r,
			capacity: capacity
		}
	}

	/* might return a new branch node if this one was split */
	fn insert_value(&mut self, tree: &mut BTree<K>, key: K, value: schema::TID) -> Option<LazyNode> {
		let mut place = 0;
		// locate the place where to insert
		for i in range(0, self.entries.len()) {
			if self.entries[i].page_id == 0 {
				continue
			}
			println!("Entry i {:?}, key {:?}", self.entries[i], self.entries[i].key);
			// TODO find place
		}
		let go_to_page = self.entries[place].page_id;
		if go_to_page == 0 {
			let lazy_node = tree.create_leaf_page();
			let new_node = lazy_node.load(tree.manager);
			let split_candidate = match new_node {
				Leaf(mut n) => n.insert_value(key, value),
				Branch(_) => fail!("Did not create a leaf page"),
			};
			self.entries[place].page_id = lazy_node.page_id;
			return split_candidate
		}

		// currently not caring about splitting branch pages
		None
	}
}

#[test]
fn simple_insert() {
	let p = Path::new(".");
	let mut manager = buffer::BufferManager::new(1024, p.clone());
	let mut bt = BTree::new(23, 42, &mut manager);
	bt.insert(42, schema::TID::new(0, 0));
	assert!(false);
}

#[test]
fn create_leaf_page() {
	let p = Path::new(".");
	let mut manager = buffer::BufferManager::new(1024, p.clone());
	let pagelock = manager.fix_page(0).unwrap();
	let page = pagelock.read();
	let lp: LeafNode<u64> = LeafNode::new(page.get_data());
	println!("lp len: {}", lp.entries.len());
	assert!(false);
}
<commit_msg>Only missing shift for insertion now<commit_after>use std::cast;
use std::mem::size_of;
use std::raw::Slice;
use std::num::Zero;

mod buffer;
mod schema;

static LEAF_MARKER: u8 = 0b111111111;
static BRANCH_MARKER: u8 = 0b0;

struct BTree<'a, K> {
	segment: u64,
	manager: &'a mut buffer::BufferManager,
	tree: LazyNode,
	next_free_page: u64,
}

impl<'a, K: TotalOrd + Zero> BTree<'a, K> {
	fn new<'b>(segment_id: u64, dummy: K, manager: &'b mut buffer::BufferManager) -> BTree<'b, K> {
		BTree {
			segment: segment_id,
			manager: manager,
			tree: LazyNode::new(1),
			next_free_page: 2,
		}
	}

	//fn locate_page

	fn insert(&mut self, key: K, value: schema::TID) {
		let node = self.tree.load(self.manager);
		// try insertion and see if the root was split
		let candidate = match node {
			Branch(mut n) => n.insert_value(self, key, value),
			Leaf(mut n) => n.insert_value(key, value),
		};
		// set new tree root if it was split
		match candidate {
			Some(new_node) => self.tree = new_node,
			None => (),
		}
	}

	fn next_page(&mut self) -> u64 {
		let n = self.next_free_page;
		self.next_free_page += 1;
		n
	}

	fn create_branch_page(&mut self) -> LazyNode {
		let next = self.next_page();
		let page_path = buffer::join_segment(self.segment, next);
		let pagelock = self.manager.fix_page(page_path).unwrap();
		let mut page = pagelock.write();
		let data = page.get_mut_data();
		data[0] = 0;
		LazyNode::new(page_path)
	}

	fn create_leaf_page(&mut self) -> LazyNode {
		let next = self.next_page();
		let page_path = buffer::join_segment(self.segment, next);
		let pagelock = self.manager.fix_page(page_path).unwrap();
		let mut page = pagelock.write();
		let data = page.get_mut_data();
		// marker to be a leaf page
		data[0] = 255;
		LazyNode::new(page_path)

	}

	fn erase(&mut self, key: K) {
	}

	fn lookup(self, key: K) -> Option<schema::TID> {
		None
	}
}

/* a placeholder for the actual page */
struct LazyNode {
	page_id: u64,
}

// TODO implement Drop to unfix page
impl LazyNode {
	fn new(page_id: u64) -> LazyNode {
		LazyNode {page_id: page_id}
	}

	/*
	 * As this is just a placeholder, return the actual node that his is
	 * representing
	 */
	fn load<'a, K: TotalOrd + Zero>(&self, manager: &mut buffer::BufferManager) -> Node<'a, K> {
		let pagelock = manager.fix_page(self.page_id).unwrap();
		let page = pagelock.read();
		LazyNode::load_node(page.get_data())
	}

	/* helper method to load a Node tepending on the page marker */
	fn load_node<'a, K: TotalOrd + Zero>(page: &[u8]) -> Node<'a, K> {
		if page[0] == 255 {
			Leaf(LeafNode::new(page))
		} else {
			Branch(BranchNode::new(page))
		}
	}
}

/* a node might either be an inner node (branch node) or a leaf node (LeafNode) */
enum Node<'a, K> {
	Branch(BranchNode<'a, K>),
	Leaf(LeafNode<'a, K>),
}

/* leaf pages consist mainly of leaf entries which are (K, TID) pairs */
struct LeafEntry<K> {
	key: K,
	tid: schema::TID,
}

/* branch pages consist mainly of branch entries which are (K, page_id) pairs */
struct BranchEntry<K> {
	key: K,
	page_id: u64,
}

struct LeafNode<'a, K> {
	capacity: uint,
	entries: &'a mut [LeafEntry<K>],
}

impl<'a, K: TotalOrd + Zero> LeafNode<'a, K> {
	fn new(page: &[u8]) -> LeafNode<'a, K> {
		// first byte is a Leaf/Branch marker
		let entry_size = size_of::<LeafEntry<K>>();
		let entry_num = (page.len() - size_of::<u8>()) / entry_size;

		let mut r: &mut[LeafEntry<K>] = unsafe {
			cast::transmute(
				Slice::<LeafEntry<K>> {
					data: page.as_ptr() as *() as *LeafEntry<K>,
					len: entry_num,
				}
			)
		};

		/* find slots that are used */
		let mut capacity = r.len();
		for i in range(0, r.len()) {
			if !r[i].key.is_zero() && !r[i].tid.is_invalid() {
				capacity -= 1;
			}
		}

		LeafNode {
			entries: r,
			capacity: capacity
		}
	}

	fn insert_value(&mut self, key: K, tid: schema::TID) -> Option<LazyNode> {
		if self.capacity == 0 {
			// TODO split
			fail!("shit");
			return None
		}

		// find place to insert
		let location = self.find_slot(&key);
		// free that spot
		self.shift_from(location);
		// and put it in
		self.entries[location].key = key;
		self.entries[location].tid = tid;

		// insertion went fine, done
		None
	}

	/* finds the location at which a key should be inserted */
	fn find_slot(&self, key: &K) -> uint {
		let mut found = 0;
		for i in range(0, self.entries.len()) {
			println!("Entry {:?}", self.entries[i]);
			if &self.entries[i].key > key {
				found = i - 1;
			}
		}
		found
	}

	/* moves all items from `index` one number back */
	fn shift_from(&mut self, index: uint) {
		// TODO
	}
}

struct BranchNode<'a, K> {
	capacity: uint,
	entries: &'a mut [BranchEntry<K>],
}

impl<'a, K: TotalOrd + Zero> BranchNode<'a, K> {
	fn new(page: &[u8]) -> BranchNode<'a, K> {
		let entry_size = size_of::<BranchEntry<K>>();
		let entry_num = (page.len() - size_of::<u8>()) / entry_size;
		let mut r: &mut [BranchEntry<K>] = unsafe {
			cast::transmute(
				Slice::<BranchEntry<K>> {
					data: page.as_ptr() as *() as *BranchEntry<K>,
					len: entry_num,
				}
			)
		};

		/* find slots that are used */
		let mut capacity = r.len();
		for i in range(0, r.len()) {
			if r[i].page_id != 0 {
				capacity -= 1;
			}
		}

		println!("BranchNode computed capacity: {}", capacity);

		BranchNode {
			entries: r,
			capacity: capacity
		}
	}

	/* might return a new branch node if this one was split */
	fn insert_value(&mut self, tree: &mut BTree<K>, key: K, value: schema::TID) -> Option<LazyNode> {
		let mut place = 0;
		// locate the place where to insert
		for i in range(0, self.entries.len()) {
			if self.entries[i].page_id == 0 {
				continue
			}
			println!("Entry i {:?}, key {:?}", self.entries[i], self.entries[i].key);
			// TODO find place
		}
		let go_to_page = self.entries[place].page_id;
		if go_to_page == 0 {
			let lazy_node = tree.create_leaf_page();
			let new_node = lazy_node.load(tree.manager);
			let split_candidate = match new_node {
				Leaf(mut n) => n.insert_value(key, value),
				Branch(_) => fail!("Did not create a leaf page"),
			};
			self.entries[place].page_id = lazy_node.page_id;
			return split_candidate
		}

		// currently not caring about splitting branch pages
		None
	}
}

#[test]
fn simple_insert() {
	let p = Path::new(".");
	let mut manager = buffer::BufferManager::new(1024, p.clone());
	let mut bt = BTree::new(23, 42, &mut manager);
	bt.insert(42, schema::TID::new(0, 0));
	assert!(false);
}

#[test]
fn create_leaf_page() {
	let p = Path::new(".");
	let mut manager = buffer::BufferManager::new(1024, p.clone());
	let pagelock = manager.fix_page(0).unwrap();
	let page = pagelock.read();
	let lp: LeafNode<u64> = LeafNode::new(page.get_data());
	println!("lp len: {}", lp.entries.len());
	assert!(false);
}
<|endoftext|>