<commit_before>//! Graphics draw state.

/// Graphics draw state used for blending, clipping and stencil rendering.
#[derive(Copy, Clone, PartialEq, Debug, PartialOrd)]
pub struct DrawState {
    /// Scissor mask to use. If set, no pixel outside of this
    /// rectangle (in screen space) will be written to as a result of rendering.
    pub scissor: Option<[u32; 4]>,
    /// Stencil test to use. If None, no stencil testing is done.
    pub stencil: Option<Stencil>,
    /// Blend function to use. If None, blending is disabled.
    pub blend: Option<Blend>,
}

impl Default for DrawState {
    fn default() -> Self {
        DrawState::new_alpha()
    }
}

impl DrawState {
    /// Uses alpha blending.
    pub fn new_alpha() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: None,
            scissor: None,
        }
    }

    /// Draws to stencil buffer with value 255.
    /// This can be used for clipping.
    pub fn new_clip() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: Some(Stencil::Clip(255)),
            scissor: None,
        }
    }

    /// Tests against stencil buffer with value 255.
    /// Draws inside the shape defined by stencil buffer.
    pub fn new_inside() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: Some(Stencil::Inside(255)),
            scissor: None,
        }
    }

    /// Tests against stencil buffer with value 255.
    /// Draws outside the shape defined by stencil buffer.
    pub fn new_outside() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: Some(Stencil::Outside(255)),
            scissor: None,
        }
    }

    /// Sets blending.
    pub fn blend(mut self, blend: Blend) -> DrawState {
        self.blend = Some(blend);
        self
    }

    /// Sets scissor `[x, y, w, h]`.
    pub fn scissor(mut self, scissor: [u32; 4]) -> DrawState {
        self.scissor = Some(scissor);
        self
    }
}

/// The blend setting to use when drawing.
///
/// Using presets since some backends need one pipeline state object instance
/// per blending technique.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum Blend {
    /// Alpha blending (allows semi-transparent pixels).
    ///
    /// ```not_rust
    /// new_dest_color = src_color * src_alpha + dest_color * (1 - src_alpha)
    /// new_dest_alpha = src_alpha + dest_alpha
    /// ```
    Alpha,
    /// Additive blending.
    ///
    /// ```not_rust
    /// new_dest_color = src_color + dest_color
    /// new_dest_alpha = src_alpha + dest_alpha
    /// ```
    Add,
    /// Multiply color components.
    ///
    /// ```not_rust
    /// new_dest_color = src_color * dest_color
    /// new_dest_alpha = src_alpha * dest_alpha
    /// ```
    Multiply,
    /// Invert colors when rendering a white shape.
    ///
    /// ```not_rust
    /// new_dest_color = ref_color - src_color
    /// new_dest_alpha = dest_alpha
    /// ```
    ///
    /// When combining two fragments, subtract the destination color from a constant color
    /// using the source color as weight. Has an invert effect with the constant color
    /// as base and source color controlling displacement from the base color.
    /// A white source color and a white value results in plain invert.
    /// The output alpha is same as destination alpha.
    Invert,
}

/// Stencil buffer settings.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum Stencil {
    /// Draw to stencil buffer.
    Clip(u8),
    /// Draw pixels that have stencil value.
    Inside(u8),
    /// Draw pixels that does not have stencil value.
    Outside(u8),
}
<commit_msg>Adding Blend::Screen<commit_after>//! Graphics draw state.

/// Graphics draw state used for blending, clipping and stencil rendering.
#[derive(Copy, Clone, PartialEq, Debug, PartialOrd)]
pub struct DrawState {
    /// Scissor mask to use. If set, no pixel outside of this
    /// rectangle (in screen space) will be written to as a result of rendering.
    pub scissor: Option<[u32; 4]>,
    /// Stencil test to use. If None, no stencil testing is done.
    pub stencil: Option<Stencil>,
    /// Blend function to use. If None, blending is disabled.
    pub blend: Option<Blend>,
}

impl Default for DrawState {
    fn default() -> Self {
        DrawState::new_alpha()
    }
}

impl DrawState {
    /// Uses alpha blending.
    pub fn new_alpha() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: None,
            scissor: None,
        }
    }

    /// Draws to stencil buffer with value 255.
    /// This can be used for clipping.
    pub fn new_clip() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: Some(Stencil::Clip(255)),
            scissor: None,
        }
    }

    /// Tests against stencil buffer with value 255.
    /// Draws inside the shape defined by stencil buffer.
    pub fn new_inside() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: Some(Stencil::Inside(255)),
            scissor: None,
        }
    }

    /// Tests against stencil buffer with value 255.
    /// Draws outside the shape defined by stencil buffer.
    pub fn new_outside() -> DrawState {
        DrawState {
            blend: Some(Blend::Alpha),
            stencil: Some(Stencil::Outside(255)),
            scissor: None,
        }
    }

    /// Sets blending.
    pub fn blend(mut self, blend: Blend) -> DrawState {
        self.blend = Some(blend);
        self
    }

    /// Sets scissor `[x, y, w, h]`.
    pub fn scissor(mut self, scissor: [u32; 4]) -> DrawState {
        self.scissor = Some(scissor);
        self
    }
}

/// The blend setting to use when drawing.
///
/// Using presets since some backends need one pipeline state object instance
/// per blending technique.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum Blend {
    /// Alpha blending (allows semi-transparent pixels).
    ///
    /// ```not_rust
    /// new_dest_color = src_color * src_alpha + dest_color * (1 - src_alpha)
    /// new_dest_alpha = src_alpha + dest_alpha
    /// ```
    Alpha,
    /// Additive blending.
    ///
    /// ```not_rust
    /// new_dest_color = src_color + dest_color
    /// new_dest_alpha = src_alpha + dest_alpha
    /// ```
    Add,
    /// Screen blending.
    ///
    /// ```not_rust
    /// new_dest_color = src_color * src_alpha + dest_color
    /// new_dest_alpha = src_alpha + dest_alpha
    /// ```
    Screen,
    /// Multiply color components.
    ///
    /// ```not_rust
    /// new_dest_color = src_color * dest_color
    /// new_dest_alpha = src_alpha * dest_alpha
    /// ```
    Multiply,
    /// Invert colors when rendering a white shape.
    ///
    /// ```not_rust
    /// new_dest_color = ref_color - src_color
    /// new_dest_alpha = dest_alpha
    /// ```
    ///
    /// When combining two fragments, subtract the destination color from a constant color
    /// using the source color as weight. Has an invert effect with the constant color
    /// as base and source color controlling displacement from the base color.
    /// A white source color and a white value results in plain invert.
    /// The output alpha is same as destination alpha.
    Invert,
}

/// Stencil buffer settings.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum Stencil {
    /// Draw to stencil buffer.
    Clip(u8),
    /// Draw pixels that have stencil value.
    Inside(u8),
    /// Draw pixels that does not have stencil value.
    Outside(u8),
}
<|endoftext|>


<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,
    );

    (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(LOCAL_CRATE);
    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 (main_def_id, entry_ty) = tcx.entry_fn(LOCAL_CRATE).unwrap();
    let instance = Instance::mono(tcx, main_def_id).polymorphize(tcx);

    match entry_ty {
        EntryFnType::Main => {
            // FIXME set program arguments somehow

            let main_sig = Signature {
                params: vec![],
                returns: vec![],
                call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
            };
            let main_func_id = jit_module
                .declare_function(tcx.symbol_name(instance).name, Linkage::Import, &main_sig)
                .unwrap();
            let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);

            CURRENT_MODULE.with(|current_module| {
                assert!(current_module.borrow_mut().replace(jit_module).is_none())
            });

            let f: extern "C" fn() = unsafe { ::std::mem::transmute(finalized_main) };
            f();
            std::process::exit(0);
        }
        EntryFnType::Start => {
            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(tcx.symbol_name(instance).name, Linkage::Import, &start_sig)
                .unwrap();
            let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);

            CURRENT_MODULE.with(|current_module| {
                assert!(current_module.borrow_mut().replace(jit_module).is_none())
            });

            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(LOCAL_CRATE);
    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>[cg_clif] Fix run_jit from sync<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 rustc_session::config::EntryFnType;

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,
    );

    (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(LOCAL_CRATE);
    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(LOCAL_CRATE);
    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();
}
<|endoftext|>


<commit_before>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines. Any line that starts with a # is ignored.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Opcode info from game packets
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcast packet type opcodes.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::thread::Builder;
use std::sync::{Arc, Mutex};

extern crate time;
use time::{SteadyTime,Duration};

mod error_macros;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::{HeartbeatPacket,GAMEEXITED,VALID_GAME_VERSIONS,HALO_RETAIL};

trait IPString {
    fn ip_string(&self) -> String;
}

impl IPString for SocketAddr {
    fn ip_string(&self) -> String {
        match *self {
            V4(ipv4) => ipv4.ip().to_string(),
            V6(ipv6) => "[".to_owned() + &ipv6.ip().to_string() + "]"
        }
    }
}

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let args : Vec<_> = env::args().collect();
        args[1].to_owned()
    }
    else if count == 1 {
        "0.0.0.0".to_owned()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = unwrap_result_or_bail!(UdpSocket::bind((&ip as &str,BROADCAST_PORT_UDP)), {
        println!("Failed creating a UDP socket at {}:{}.",ip,BROADCAST_PORT_UDP);
        return;
    });

    let client_socket = unwrap_result_or_bail!(TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)), {
        println!("Failed listening to TCP at {}:{}.",ip,SERVER_LIST_PORT_TCP);
        return;
    });

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    let _ = Builder::new().name("destruction_thread".to_owned()).spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let timenow = SteadyTime::now();
            servers.retain(|x| x.last_alive + Duration::seconds(DROP_TIME) > timenow);
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Option<Vec<String>> = None;
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    let _ = Builder::new().name("blacklist_thread".to_owned()).spawn(move || {
        let valid_line = |x: &str| -> bool { x.trim().len() > 0 && !x.starts_with("#") };
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_ref = blacklist_update.lock().unwrap();
                *blacklist_ref =
                    File::open(BLACKLIST_FILE).
                    map(|file|
                        BufReader::new(&file).lines().
                        filter_map(|line| line.ok().and_then(|x| if valid_line(&x) { Some(x) } else { None })).
                        collect()
                    ).ok();
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    let _ = Builder::new().name("halomd_thread".to_owned()).spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = unwrap_option_or_bail!(stream.ok(), { continue });
                let ip = unwrap_option_or_bail!(client.peer_addr().ok(), { continue }).ip_string();
                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips.push_str(&(j.to_string()));
                        ips.push('\n');
                    }
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips.push_str(&ip);
                ips.push_str(":49149:3425");

                // We may be here a while. Just in case...
                // May want to consider some lightweight thread library in the future, or some non-blocking polling mechanism
                thread::spawn( move || {
                    let _ = client.write_all(ips.as_bytes());
                });
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = unwrap_option_or_bail!(halo_socket.recv_from(&mut buffer).ok(), { continue });

        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = source.ip_string();

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if (*blacklist_ref).as_ref().map(|blacklist| blacklist.contains(&client_ip)).unwrap_or(false) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {
            let mut servers = servers_mut_udp.lock().unwrap();

            match HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]) {
                None => {},
                Some(packet) => {
                    let updatetime = SteadyTime::now();
                    match servers.iter_mut().position(|x| *x == (&client_ip, packet.localport)) {
                        None => {
                            if packet.gamename == HALO_RETAIL && VALID_GAME_VERSIONS.contains(&&*packet.gamever) {
                                (*servers).push(HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime });
                            }
                        }
                        Some(k) => {
                            servers[k].last_alive = updatetime;
                            if packet.statechanged == GAMEEXITED {
                                servers.remove(k);
                            }
                        }
                    };
                }
            };
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            let mut servers_ref = servers_mut_udp.lock().unwrap();
            let servers = (*servers_ref).iter_mut();

            for i in servers {
                if *i == (&client_ip, source.port()) {
                    i.last_alive = SteadyTime::now();
                    break;
                }
            }
        }
    }
}
<commit_msg>Make constants for thread names<commit_after>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines. Any line that starts with a # is ignored.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Opcode info from game packets
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcast packet type opcodes.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

const TCP_SERVER_THREAD_NAME : &'static str = "halomd_thread";
const DESTRUCTION_THREAD_NAME : &'static str = "destruction_thread";
const BLACKLIST_THREAD_NAME : &'static str = "blacklist_thread";

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::thread::Builder;
use std::sync::{Arc, Mutex};

extern crate time;
use time::{SteadyTime,Duration};

mod error_macros;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::{HeartbeatPacket,GAMEEXITED,VALID_GAME_VERSIONS,HALO_RETAIL};

trait IPString {
    fn ip_string(&self) -> String;
}

impl IPString for SocketAddr {
    fn ip_string(&self) -> String {
        match *self {
            V4(ipv4) => ipv4.ip().to_string(),
            V6(ipv6) => "[".to_owned() + &ipv6.ip().to_string() + "]"
        }
    }
}

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let args : Vec<_> = env::args().collect();
        args[1].to_owned()
    }
    else if count == 1 {
        "0.0.0.0".to_owned()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = unwrap_result_or_bail!(UdpSocket::bind((&ip as &str,BROADCAST_PORT_UDP)), {
        println!("Failed creating a UDP socket at {}:{}.",ip,BROADCAST_PORT_UDP);
        return;
    });

    let client_socket = unwrap_result_or_bail!(TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)), {
        println!("Failed listening to TCP at {}:{}.",ip,SERVER_LIST_PORT_TCP);
        return;
    });

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    let _ = Builder::new().name(DESTRUCTION_THREAD_NAME.to_owned()).spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let timenow = SteadyTime::now();
            servers.retain(|x| x.last_alive + Duration::seconds(DROP_TIME) > timenow);
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Option<Vec<String>> = None;
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    let _ = Builder::new().name(BLACKLIST_THREAD_NAME.to_owned()).spawn(move || {
        let valid_line = |x: &str| -> bool { x.trim().len() > 0 && !x.starts_with("#") };
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_ref = blacklist_update.lock().unwrap();
                *blacklist_ref =
                    File::open(BLACKLIST_FILE).
                    map(|file|
                        BufReader::new(&file).lines().
                        filter_map(|line| line.ok().and_then(|x| if valid_line(&x) { Some(x) } else { None })).
                        collect()
                    ).ok();
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    let _ = Builder::new().name(TCP_SERVER_THREAD_NAME.to_owned()).spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = unwrap_option_or_bail!(stream.ok(), { continue });
                let ip = unwrap_option_or_bail!(client.peer_addr().ok(), { continue }).ip_string();
                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips.push_str(&(j.to_string()));
                        ips.push('\n');
                    }
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips.push_str(&ip);
                ips.push_str(":49149:3425");

                // We may be here a while. Just in case...
                // May want to consider some lightweight thread library in the future, or some non-blocking polling mechanism
                thread::spawn( move || {
                    let _ = client.write_all(ips.as_bytes());
                });
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = unwrap_option_or_bail!(halo_socket.recv_from(&mut buffer).ok(), { continue });

        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = source.ip_string();

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if (*blacklist_ref).as_ref().map(|blacklist| blacklist.contains(&client_ip)).unwrap_or(false) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {
            let mut servers = servers_mut_udp.lock().unwrap();

            match HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]) {
                None => {},
                Some(packet) => {
                    let updatetime = SteadyTime::now();
                    match servers.iter_mut().position(|x| *x == (&client_ip, packet.localport)) {
                        None => {
                            if packet.gamename == HALO_RETAIL && VALID_GAME_VERSIONS.contains(&&*packet.gamever) {
                                (*servers).push(HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime });
                            }
                        }
                        Some(k) => {
                            servers[k].last_alive = updatetime;
                            if packet.statechanged == GAMEEXITED {
                                servers.remove(k);
                            }
                        }
                    };
                }
            };
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            let mut servers_ref = servers_mut_udp.lock().unwrap();
            let servers = (*servers_ref).iter_mut();

            for i in servers {
                if *i == (&client_ip, source.port()) {
                    i.last_alive = SteadyTime::now();
                    break;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use hir::db::HirDatabase;
use ra_syntax::{algo::non_trivia_sibling, Direction, T};

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn flip_comma(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let comma = ctx.token_at_offset().find(|leaf| leaf.kind() == T![,])?;
    let prev = non_trivia_sibling(comma.clone().into(), Direction::Prev)?;
    let next = non_trivia_sibling(comma.clone().into(), Direction::Next)?;
    ctx.add_action(AssistId("flip_comma"), "flip comma", |edit| {
        edit.target(comma.text_range());
        edit.replace(prev.text_range(), next.to_string());
        edit.replace(next.text_range(), prev.to_string());
    });

    ctx.build()
}

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

    use crate::helpers::{check_assist, check_assist_target};

    #[test]
    fn flip_comma_works_for_function_parameters() {
        check_assist(
            flip_comma,
            "fn foo(x: i32,<|> y: Result<(), ()>) {}",
            "fn foo(y: Result<(), ()>,<|> x: i32) {}",
        )
    }

    #[test]
    fn flip_comma_target() {
        check_assist_target(flip_comma, "fn foo(x: i32,<|> y: Result<(), ()>) {}", ",")
    }
}
<commit_msg>Fix flip comma assist<commit_after>use hir::db::HirDatabase;
use ra_syntax::{algo::non_trivia_sibling, Direction, T};

use crate::{Assist, AssistCtx, AssistId};

pub(crate) fn flip_comma(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let comma = ctx.token_at_offset().find(|leaf| leaf.kind() == T![,])?;
    let prev = non_trivia_sibling(comma.clone().into(), Direction::Prev)?;
    let next = non_trivia_sibling(comma.clone().into(), Direction::Next)?;

    // Don't apply a "flip" in case of a last comma
    // that typically comes before punctuation
    if next.kind().is_punct() {
        return None;
    }

    ctx.add_action(AssistId("flip_comma"), "flip comma", |edit| {
        edit.target(comma.text_range());
        edit.replace(prev.text_range(), next.to_string());
        edit.replace(next.text_range(), prev.to_string());
    });

    ctx.build()
}

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

    use crate::helpers::{check_assist, check_assist_target};

    #[test]
    fn flip_comma_works_for_function_parameters() {
        check_assist(
            flip_comma,
            "fn foo(x: i32,<|> y: Result<(), ()>) {}",
            "fn foo(y: Result<(), ()>,<|> x: i32) {}",
        )
    }

    #[test]
    fn flip_comma_target() {
        check_assist_target(flip_comma, "fn foo(x: i32,<|> y: Result<(), ()>) {}", ",")
    }

    #[test]
    #[should_panic]
    fn flip_comma_before_punct() {
        // See https://github.com/rust-analyzer/rust-analyzer/issues/1619
        // "Flip comma" assist shouldn't be applicable to the last comma in enum or struct
        // declaration body.
        check_assist_target(flip_comma,
                            "pub enum Test { \
                                A,<|> \
                            }", ",");

        check_assist_target(flip_comma,
                            "pub struct Test { \
                                foo: usize,<|> \
                            }", ",");
    }
}
<|endoftext|>


<commit_before>#![cfg(feature = "quickcheck")]
#![cfg(feature = "quickcheck_macros")]

extern crate num_bigint;
extern crate num_integer;
extern crate num_traits;

extern crate quickcheck;
#[macro_use]
extern crate quickcheck_macros;

use num_bigint::{BigInt, BigUint};
use num_traits::{Num, One, Pow, Zero};
use quickcheck::{QuickCheck, StdThreadGen, TestResult};

#[quickcheck]
fn quickcheck_unsigned_eq_reflexive(a: BigUint) -> bool {
    a == a
}

#[quickcheck]
fn quickcheck_signed_eq_reflexive(a: BigInt) -> bool {
    a == a
}

#[quickcheck]
fn quickcheck_unsigned_eq_symmetric(a: BigUint, b: BigUint) -> bool {
    if a == b {
        b == a
    } else {
        b != a
    }
}

#[quickcheck]
fn quickcheck_signed_eq_symmetric(a: BigInt, b: BigInt) -> bool {
    if a == b {
        b == a
    } else {
        b != a
    }
}

#[test]
fn quickcheck_arith_primitive() {
    let gen = StdThreadGen::new(usize::max_value());
    let mut qc = QuickCheck::with_gen(gen);

    fn test_unsigned_add_primitive(a: usize, b: usize) -> TestResult {
        let actual = BigUint::from(a) + BigUint::from(b);
        match a.checked_add(b) {
            None => TestResult::discard(),
            Some(expected) => TestResult::from_bool(BigUint::from(expected) == actual),
        }
    }

    fn test_signed_add_primitive(a: isize, b: isize) -> TestResult {
        let actual = BigInt::from(a) + BigInt::from(b);
        match a.checked_add(b) {
            None => TestResult::discard(),
            Some(expected) => TestResult::from_bool(BigInt::from(expected) == actual),
        }
    }

    fn test_unsigned_mul_primitive(a: u64, b: u64) -> bool {
        //maximum value of u64 means no overflow
        BigUint::from(a as u128 * b as u128) == BigUint::from(a) * BigUint::from(b)
    }

    fn test_signed_mul_primitive(a: i64, b: i64) -> bool {
        //maximum value of i64 means no overflow
        BigInt::from(a as i128 * b as i128) == BigInt::from(a) * BigInt::from(b)
    }

    fn test_unsigned_sub_primitive(a: u128, b: u128) -> bool {
        if b < a {
            BigUint::from(a - b) == BigUint::from(a) - BigUint::from(b)
        } else {
            BigUint::from(b - a) == BigUint::from(b) - BigUint::from(a)
        }
    }

    fn test_signed_sub_primitive(a: i128, b: i128) -> bool {
        if b < a {
            BigInt::from(a - b) == BigInt::from(a) - BigInt::from(b)
        } else {
            BigInt::from(b - a) == BigInt::from(b) - BigInt::from(a)
        }
    }

    fn test_unsigned_div_primitive(a: u128, b: u128) -> TestResult {
        if b == 0 {
            TestResult::discard()
        } else {
            TestResult::from_bool(BigUint::from(a / b) == BigUint::from(a) / BigUint::from(b))
        }
    }

    fn test_signed_div_primitive(a: i128, b: i128) -> TestResult {
        if b == 0 {
            TestResult::discard()
        } else {
            TestResult::from_bool(BigInt::from(a / b) == BigInt::from(a) / BigInt::from(b))
        }
    }

    qc.quickcheck(test_unsigned_add_primitive as fn(usize, usize) -> TestResult);
    qc.quickcheck(test_signed_add_primitive as fn(isize, isize) -> TestResult);
    qc.quickcheck(test_unsigned_mul_primitive as fn(u64, u64) -> bool);
    qc.quickcheck(test_signed_mul_primitive as fn(i64, i64) -> bool);
    qc.quickcheck(test_unsigned_sub_primitive as fn(u128, u128) -> bool);
    qc.quickcheck(test_signed_sub_primitive as fn(i128, i128) -> bool);
    qc.quickcheck(test_unsigned_div_primitive as fn(u128, u128) -> TestResult);
    qc.quickcheck(test_signed_div_primitive as fn(i128, i128) -> TestResult);
}

#[quickcheck]
fn quickcheck_unsigned_add_commutative(a: BigUint, b: BigUint) -> bool {
    &a + &b == b + a
}

#[quickcheck]
fn quickcheck_signed_add_commutative(a: BigInt, b: BigInt) -> bool {
    &a + &b == b + a
}

#[quickcheck]
fn quickcheck_unsigned_add_zero(a: BigUint) -> bool {
    a == &a + BigUint::zero()
}

#[quickcheck]
fn quickcheck_signed_add_zero(a: BigInt) -> bool {
    a == &a + BigInt::zero()
}

#[quickcheck]
fn quickcheck_unsigned_add_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
    (&a + &b) + &c == a + (b + c)
}

#[quickcheck]
fn quickcheck_signed_add_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
    (&a + &b) + &c == a + (b + c)
}

#[quickcheck]
fn quickcheck_unsigned_mul_zero(a: BigUint) -> bool {
    a * BigUint::zero() == BigUint::zero()
}

#[quickcheck]
fn quickcheck_signed_mul_zero(a: BigInt) -> bool {
    a * BigInt::zero() == BigInt::zero()
}

#[quickcheck]
fn quickcheck_unsigned_mul_one(a: BigUint) -> bool {
    &a * BigUint::one() == a
}

#[quickcheck]
fn quickcheck_signed_mul_one(a: BigInt) -> bool {
    &a * BigInt::one() == a
}

#[quickcheck]
fn quickcheck_unsigned_mul_commutative(a: BigUint, b: BigUint) -> bool {
    &a * &b == b * a
}

#[quickcheck]
fn quickcheck_signed_mul_commutative(a: BigInt, b: BigInt) -> bool {
    &a * &b == b * a
}

#[quickcheck]
fn quickcheck_unsigned_mul_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
    (&a * &b) * &c == a * (b * c)
}

#[quickcheck]
fn quickcheck_signed_mul_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
    (&a * &b) * &c == a * (b * c)
}

#[quickcheck]
fn quickcheck_unsigned_distributive(a: BigUint, b: BigUint, c: BigUint) -> bool {
    &a * (&b + &c) == &a * b + a * c
}

#[quickcheck]
fn quickcheck_signed_distributive(a: BigInt, b: BigInt, c: BigInt) -> bool {
    &a * (&b + &c) == &a * b + a * c
}

#[quickcheck]
///Tests that exactly one of a<b a>b a=b is true
fn quickcheck_unsigned_ge_le_eq_mut_exclusive(a: BigUint, b: BigUint) -> bool {
    let gt_lt_eq = vec![a > b, a < b, a == b];
    gt_lt_eq
        .iter()
        .fold(0, |acc, e| if *e { acc + 1 } else { acc })
        == 1
}

#[quickcheck]
///Tests that exactly one of a<b a>b a=b is true
fn quickcheck_signed_ge_le_eq_mut_exclusive(a: BigInt, b: BigInt) -> bool {
    let gt_lt_eq = vec![a > b, a < b, a == b];
    gt_lt_eq
        .iter()
        .fold(0, |acc, e| if *e { acc + 1 } else { acc })
        == 1
}

#[quickcheck]
/// Tests correctness of subtraction assuming addition is correct
fn quickcheck_unsigned_sub(a: BigUint, b: BigUint) -> bool {
    if b < a {
        &a - &b + b == a
    } else {
        &b - &a + a == b
    }
}

#[quickcheck]
/// Tests correctness of subtraction assuming addition is correct
fn quickcheck_signed_sub(a: BigInt, b: BigInt) -> bool {
    if b < a {
        &a - &b + b == a
    } else {
        &b - &a + a == b
    }
}

#[quickcheck]
fn quickcheck_unsigned_pow_zero(a: BigUint) -> bool {
    a.pow(0_u32) == BigUint::one()
}

#[quickcheck]
fn quickcheck_unsigned_pow_one(a: BigUint) -> bool {
    a.pow(1_u32) == a
}

#[quickcheck]
fn quickcheck_unsigned_sqrt(a: BigUint) -> bool {
    (&a * &a).sqrt() == a
}

#[quickcheck]
fn quickcheck_unsigned_cbrt(a: BigUint) -> bool {
    (&a * &a * &a).cbrt() == a
}

#[quickcheck]
fn quickcheck_signed_cbrt(a: BigInt) -> bool {
    (&a * &a * &a).cbrt() == a
}

#[quickcheck]
fn quickcheck_unsigned_conversion(a: BigUint, radix: u8) -> TestResult {
    let radix = radix as u32;
    if radix > 36 || radix < 2 {
        return TestResult::discard();
    }
    let string = a.to_str_radix(radix);
    TestResult::from_bool(a == BigUint::from_str_radix(&string, radix).unwrap())
}

#[quickcheck]
fn quickcheck_signed_conversion(a: BigInt, radix: u8) -> TestResult {
    let radix = radix as u32;
    if radix > 36 || radix < 2 {
        return TestResult::discard();
    }
    let string = a.to_str_radix(radix);
    TestResult::from_bool(a == BigInt::from_str_radix(&string, radix).unwrap())
}

#[test]
fn quicktest_shift() {
    let gen = StdThreadGen::new(usize::max_value());
    let mut qc = QuickCheck::with_gen(gen);

    fn test_shr_unsigned(a: u64, shift: u8) -> TestResult {
        let shift = (shift % 64) as usize; //shift at most 64 bits
        let big_a = BigUint::from(a);
        TestResult::from_bool(BigUint::from(a >> shift) == big_a >> shift)
    }

    fn test_shr_signed(a: i64, shift: u8) -> TestResult {
        let shift = (shift % 64) as usize; //shift at most 64 bits
        let big_a = BigInt::from(a);
        TestResult::from_bool(BigInt::from(a >> shift) == big_a >> shift)
    }

    fn test_shl_unsigned(a: u32, shift: u8) -> TestResult {
        let shift = (shift % 32) as usize; //shift at most 32 bits
        let a = a as u64; //leave room for the shifted bits
        let big_a = BigUint::from(a);
        TestResult::from_bool(BigUint::from(a >> shift) == big_a >> shift)
    }

    fn test_shl_signed(a: i32, shift: u8) -> TestResult {
        let shift = (shift % 32) as usize;
        let a = a as u64; //leave room for the shifted bits
        let big_a = BigInt::from(a);
        TestResult::from_bool(BigInt::from(a >> shift) == big_a >> shift)
    }

    qc.quickcheck(test_shr_unsigned as fn(u64, u8) -> TestResult);
    qc.quickcheck(test_shr_signed as fn(i64, u8) -> TestResult);
    qc.quickcheck(test_shl_unsigned as fn(u32, u8) -> TestResult);
    qc.quickcheck(test_shl_signed as fn(i32, u8) -> TestResult);
}
<commit_msg>Add a quickcheck modpow test<commit_after>#![cfg(feature = "quickcheck")]
#![cfg(feature = "quickcheck_macros")]

extern crate num_bigint;
extern crate num_integer;
extern crate num_traits;

extern crate quickcheck;
#[macro_use]
extern crate quickcheck_macros;

use num_bigint::{BigInt, BigUint};
use num_integer::Integer;
use num_traits::{Num, One, Pow, Signed, Zero};
use quickcheck::{QuickCheck, StdThreadGen, TestResult};

#[quickcheck]
fn quickcheck_unsigned_eq_reflexive(a: BigUint) -> bool {
    a == a
}

#[quickcheck]
fn quickcheck_signed_eq_reflexive(a: BigInt) -> bool {
    a == a
}

#[quickcheck]
fn quickcheck_unsigned_eq_symmetric(a: BigUint, b: BigUint) -> bool {
    if a == b {
        b == a
    } else {
        b != a
    }
}

#[quickcheck]
fn quickcheck_signed_eq_symmetric(a: BigInt, b: BigInt) -> bool {
    if a == b {
        b == a
    } else {
        b != a
    }
}

#[test]
fn quickcheck_arith_primitive() {
    let gen = StdThreadGen::new(usize::max_value());
    let mut qc = QuickCheck::with_gen(gen);

    fn test_unsigned_add_primitive(a: usize, b: usize) -> TestResult {
        let actual = BigUint::from(a) + BigUint::from(b);
        match a.checked_add(b) {
            None => TestResult::discard(),
            Some(expected) => TestResult::from_bool(BigUint::from(expected) == actual),
        }
    }

    fn test_signed_add_primitive(a: isize, b: isize) -> TestResult {
        let actual = BigInt::from(a) + BigInt::from(b);
        match a.checked_add(b) {
            None => TestResult::discard(),
            Some(expected) => TestResult::from_bool(BigInt::from(expected) == actual),
        }
    }

    fn test_unsigned_mul_primitive(a: u64, b: u64) -> bool {
        //maximum value of u64 means no overflow
        BigUint::from(a as u128 * b as u128) == BigUint::from(a) * BigUint::from(b)
    }

    fn test_signed_mul_primitive(a: i64, b: i64) -> bool {
        //maximum value of i64 means no overflow
        BigInt::from(a as i128 * b as i128) == BigInt::from(a) * BigInt::from(b)
    }

    fn test_unsigned_sub_primitive(a: u128, b: u128) -> bool {
        if b < a {
            BigUint::from(a - b) == BigUint::from(a) - BigUint::from(b)
        } else {
            BigUint::from(b - a) == BigUint::from(b) - BigUint::from(a)
        }
    }

    fn test_signed_sub_primitive(a: i128, b: i128) -> bool {
        if b < a {
            BigInt::from(a - b) == BigInt::from(a) - BigInt::from(b)
        } else {
            BigInt::from(b - a) == BigInt::from(b) - BigInt::from(a)
        }
    }

    fn test_unsigned_div_primitive(a: u128, b: u128) -> TestResult {
        if b == 0 {
            TestResult::discard()
        } else {
            TestResult::from_bool(BigUint::from(a / b) == BigUint::from(a) / BigUint::from(b))
        }
    }

    fn test_signed_div_primitive(a: i128, b: i128) -> TestResult {
        if b == 0 {
            TestResult::discard()
        } else {
            TestResult::from_bool(BigInt::from(a / b) == BigInt::from(a) / BigInt::from(b))
        }
    }

    qc.quickcheck(test_unsigned_add_primitive as fn(usize, usize) -> TestResult);
    qc.quickcheck(test_signed_add_primitive as fn(isize, isize) -> TestResult);
    qc.quickcheck(test_unsigned_mul_primitive as fn(u64, u64) -> bool);
    qc.quickcheck(test_signed_mul_primitive as fn(i64, i64) -> bool);
    qc.quickcheck(test_unsigned_sub_primitive as fn(u128, u128) -> bool);
    qc.quickcheck(test_signed_sub_primitive as fn(i128, i128) -> bool);
    qc.quickcheck(test_unsigned_div_primitive as fn(u128, u128) -> TestResult);
    qc.quickcheck(test_signed_div_primitive as fn(i128, i128) -> TestResult);
}

#[quickcheck]
fn quickcheck_unsigned_add_commutative(a: BigUint, b: BigUint) -> bool {
    &a + &b == b + a
}

#[quickcheck]
fn quickcheck_signed_add_commutative(a: BigInt, b: BigInt) -> bool {
    &a + &b == b + a
}

#[quickcheck]
fn quickcheck_unsigned_add_zero(a: BigUint) -> bool {
    a == &a + BigUint::zero()
}

#[quickcheck]
fn quickcheck_signed_add_zero(a: BigInt) -> bool {
    a == &a + BigInt::zero()
}

#[quickcheck]
fn quickcheck_unsigned_add_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
    (&a + &b) + &c == a + (b + c)
}

#[quickcheck]
fn quickcheck_signed_add_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
    (&a + &b) + &c == a + (b + c)
}

#[quickcheck]
fn quickcheck_unsigned_mul_zero(a: BigUint) -> bool {
    a * BigUint::zero() == BigUint::zero()
}

#[quickcheck]
fn quickcheck_signed_mul_zero(a: BigInt) -> bool {
    a * BigInt::zero() == BigInt::zero()
}

#[quickcheck]
fn quickcheck_unsigned_mul_one(a: BigUint) -> bool {
    &a * BigUint::one() == a
}

#[quickcheck]
fn quickcheck_signed_mul_one(a: BigInt) -> bool {
    &a * BigInt::one() == a
}

#[quickcheck]
fn quickcheck_unsigned_mul_commutative(a: BigUint, b: BigUint) -> bool {
    &a * &b == b * a
}

#[quickcheck]
fn quickcheck_signed_mul_commutative(a: BigInt, b: BigInt) -> bool {
    &a * &b == b * a
}

#[quickcheck]
fn quickcheck_unsigned_mul_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
    (&a * &b) * &c == a * (b * c)
}

#[quickcheck]
fn quickcheck_signed_mul_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
    (&a * &b) * &c == a * (b * c)
}

#[quickcheck]
fn quickcheck_unsigned_distributive(a: BigUint, b: BigUint, c: BigUint) -> bool {
    &a * (&b + &c) == &a * b + a * c
}

#[quickcheck]
fn quickcheck_signed_distributive(a: BigInt, b: BigInt, c: BigInt) -> bool {
    &a * (&b + &c) == &a * b + a * c
}

#[quickcheck]
///Tests that exactly one of a<b a>b a=b is true
fn quickcheck_unsigned_ge_le_eq_mut_exclusive(a: BigUint, b: BigUint) -> bool {
    let gt_lt_eq = vec![a > b, a < b, a == b];
    gt_lt_eq
        .iter()
        .fold(0, |acc, e| if *e { acc + 1 } else { acc })
        == 1
}

#[quickcheck]
///Tests that exactly one of a<b a>b a=b is true
fn quickcheck_signed_ge_le_eq_mut_exclusive(a: BigInt, b: BigInt) -> bool {
    let gt_lt_eq = vec![a > b, a < b, a == b];
    gt_lt_eq
        .iter()
        .fold(0, |acc, e| if *e { acc + 1 } else { acc })
        == 1
}

#[quickcheck]
/// Tests correctness of subtraction assuming addition is correct
fn quickcheck_unsigned_sub(a: BigUint, b: BigUint) -> bool {
    if b < a {
        &a - &b + b == a
    } else {
        &b - &a + a == b
    }
}

#[quickcheck]
/// Tests correctness of subtraction assuming addition is correct
fn quickcheck_signed_sub(a: BigInt, b: BigInt) -> bool {
    if b < a {
        &a - &b + b == a
    } else {
        &b - &a + a == b
    }
}

#[quickcheck]
fn quickcheck_unsigned_pow_zero(a: BigUint) -> bool {
    a.pow(0_u32) == BigUint::one()
}

#[quickcheck]
fn quickcheck_unsigned_pow_one(a: BigUint) -> bool {
    a.pow(1_u32) == a
}

#[quickcheck]
fn quickcheck_unsigned_sqrt(a: BigUint) -> bool {
    (&a * &a).sqrt() == a
}

#[quickcheck]
fn quickcheck_unsigned_cbrt(a: BigUint) -> bool {
    (&a * &a * &a).cbrt() == a
}

#[quickcheck]
fn quickcheck_signed_cbrt(a: BigInt) -> bool {
    (&a * &a * &a).cbrt() == a
}

#[quickcheck]
fn quickcheck_unsigned_conversion(a: BigUint, radix: u8) -> TestResult {
    let radix = radix as u32;
    if radix > 36 || radix < 2 {
        return TestResult::discard();
    }
    let string = a.to_str_radix(radix);
    TestResult::from_bool(a == BigUint::from_str_radix(&string, radix).unwrap())
}

#[quickcheck]
fn quickcheck_signed_conversion(a: BigInt, radix: u8) -> TestResult {
    let radix = radix as u32;
    if radix > 36 || radix < 2 {
        return TestResult::discard();
    }
    let string = a.to_str_radix(radix);
    TestResult::from_bool(a == BigInt::from_str_radix(&string, radix).unwrap())
}

#[test]
fn quicktest_shift() {
    let gen = StdThreadGen::new(usize::max_value());
    let mut qc = QuickCheck::with_gen(gen);

    fn test_shr_unsigned(a: u64, shift: u8) -> TestResult {
        let shift = (shift % 64) as usize; //shift at most 64 bits
        let big_a = BigUint::from(a);
        TestResult::from_bool(BigUint::from(a >> shift) == big_a >> shift)
    }

    fn test_shr_signed(a: i64, shift: u8) -> TestResult {
        let shift = (shift % 64) as usize; //shift at most 64 bits
        let big_a = BigInt::from(a);
        TestResult::from_bool(BigInt::from(a >> shift) == big_a >> shift)
    }

    fn test_shl_unsigned(a: u32, shift: u8) -> TestResult {
        let shift = (shift % 32) as usize; //shift at most 32 bits
        let a = a as u64; //leave room for the shifted bits
        let big_a = BigUint::from(a);
        TestResult::from_bool(BigUint::from(a >> shift) == big_a >> shift)
    }

    fn test_shl_signed(a: i32, shift: u8) -> TestResult {
        let shift = (shift % 32) as usize;
        let a = a as u64; //leave room for the shifted bits
        let big_a = BigInt::from(a);
        TestResult::from_bool(BigInt::from(a >> shift) == big_a >> shift)
    }

    qc.quickcheck(test_shr_unsigned as fn(u64, u8) -> TestResult);
    qc.quickcheck(test_shr_signed as fn(i64, u8) -> TestResult);
    qc.quickcheck(test_shl_unsigned as fn(u32, u8) -> TestResult);
    qc.quickcheck(test_shl_signed as fn(i32, u8) -> TestResult);
}

#[test]
fn quickcheck_modpow() {
    let gen = StdThreadGen::new(usize::max_value());
    let mut qc = QuickCheck::with_gen(gen);

    fn simple_modpow(base: &BigInt, exponent: &BigInt, modulus: &BigInt) -> BigInt {
        assert!(!exponent.is_negative());
        let mut result = BigInt::one().mod_floor(modulus);
        let mut base = base.mod_floor(modulus);
        let mut exponent = exponent.clone();
        while !exponent.is_zero() {
            if exponent.is_odd() {
                result = (result * &base).mod_floor(modulus);
            }
            base = (&base * &base).mod_floor(modulus);
            exponent >>= 1;
        }
        result
    }

    fn test_modpow(base: i128, exponent: u128, modulus: i128) -> TestResult {
        if modulus.is_zero() {
            TestResult::discard()
        } else {
            let base = BigInt::from(base);
            let exponent = BigInt::from(exponent);
            let modulus = BigInt::from(modulus);
            let modpow = base.modpow(&exponent, &modulus);
            let simple = simple_modpow(&base, &exponent, &modulus);
            if modpow != simple {
                eprintln!("{}.modpow({}, {})", base, exponent, modulus);
                eprintln!("  expected {}", simple);
                eprintln!("    actual {}", modpow);
                TestResult::failed()
            } else {
                TestResult::passed()
            }
        }
    }

    qc.quickcheck(test_modpow as fn(i128, u128, i128) -> TestResult);
}
<|endoftext|>


<commit_before>use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use rustc_data_structures::stack::ensure_sufficient_stack;
use rustc_hir::def_id::{DefId, LocalDefId};
use rustc_middle::mir::TerminatorKind;
use rustc_middle::ty::TypeFoldable;
use rustc_middle::ty::{self, subst::SubstsRef, InstanceDef, TyCtxt};

// FIXME: check whether it is cheaper to precompute the entire call graph instead of invoking
// this query riddiculously often.
#[instrument(level = "debug", skip(tcx, root, target))]
crate fn mir_callgraph_reachable(
    tcx: TyCtxt<'tcx>,
    (root, target): (ty::Instance<'tcx>, LocalDefId),
) -> bool {
    trace!(%root, target = %tcx.def_path_str(target.to_def_id()));
    let param_env = tcx.param_env_reveal_all_normalized(target);
    assert_ne!(
        root.def_id().expect_local(),
        target,
        "you should not call `mir_callgraph_reachable` on immediate self recursion"
    );
    assert!(
        matches!(root.def, InstanceDef::Item(_)),
        "you should not call `mir_callgraph_reachable` on shims"
    );
    assert!(
        !tcx.is_constructor(root.def_id()),
        "you should not call `mir_callgraph_reachable` on enum/struct constructor functions"
    );
    #[instrument(
        level = "debug",
        skip(tcx, param_env, target, stack, seen, recursion_limiter, caller)
    )]
    fn process(
        tcx: TyCtxt<'tcx>,
        param_env: ty::ParamEnv<'tcx>,
        caller: ty::Instance<'tcx>,
        target: LocalDefId,
        stack: &mut Vec<ty::Instance<'tcx>>,
        seen: &mut FxHashSet<ty::Instance<'tcx>>,
        recursion_limiter: &mut FxHashMap<DefId, usize>,
    ) -> bool {
        trace!(%caller);
        for &(callee, substs) in tcx.mir_inliner_callees(caller.def) {
            let substs = caller.subst_mir_and_normalize_erasing_regions(tcx, param_env, substs);
            let callee = match ty::Instance::resolve(tcx, param_env, callee, substs).unwrap() {
                Some(callee) => callee,
                None => {
                    trace!(?callee, "cannot resolve, skipping");
                    continue;
                }
            };

            // Found a path.
            if callee.def_id() == target.to_def_id() {
                return true;
            }

            if tcx.is_constructor(callee.def_id()) {
                trace!("constructors always have MIR");
                // Constructor functions cannot cause a query cycle.
                continue;
            }

            match callee.def {
                InstanceDef::Item(_) => {
                    // If there is no MIR available (either because it was not in metadata or
                    // because it has no MIR because it's an extern function), then the inliner
                    // won't cause cycles on this.
                    if !tcx.is_mir_available(callee.def_id()) {
                        trace!(?callee, "no mir available, skipping");
                        continue;
                    }
                }
                // These have no own callable MIR.
                InstanceDef::Intrinsic(_) | InstanceDef::Virtual(..) => continue,
                // These have MIR and if that MIR is inlined, substituted and then inlining is run
                // again, a function item can end up getting inlined. Thus we'll be able to cause
                // a cycle that way
                InstanceDef::VtableShim(_)
                | InstanceDef::ReifyShim(_)
                | InstanceDef::FnPtrShim(..)
                | InstanceDef::ClosureOnceShim { .. }
                | InstanceDef::CloneShim(..) => {}
                InstanceDef::DropGlue(..) => {
                    // FIXME: A not fully substituted drop shim can cause ICEs if one attempts to
                    // have its MIR built. Likely oli-obk just screwed up the `ParamEnv`s, so this
                    // needs some more analysis.
                    if callee.needs_subst() {
                        continue;
                    }
                }
            }

            if seen.insert(callee) {
                let recursion = recursion_limiter.entry(callee.def_id()).or_default();
                trace!(?callee, recursion = *recursion);
                if tcx.sess.recursion_limit().value_within_limit(*recursion) {
                    *recursion += 1;
                    stack.push(callee);
                    let found_recursion = ensure_sufficient_stack(|| {
                        process(tcx, param_env, callee, target, stack, seen, recursion_limiter)
                    });
                    if found_recursion {
                        return true;
                    }
                    stack.pop();
                } else {
                    // Pessimistically assume that there could be recursion.
                    return true;
                }
            }
        }
        false
    }
    process(
        tcx,
        param_env,
        root,
        target,
        &mut Vec::new(),
        &mut FxHashSet::default(),
        &mut FxHashMap::default(),
    )
}

crate fn mir_inliner_callees<'tcx>(
    tcx: TyCtxt<'tcx>,
    instance: ty::InstanceDef<'tcx>,
) -> &'tcx [(DefId, SubstsRef<'tcx>)] {
    let steal;
    let guard;
    let body = match (instance, instance.def_id().as_local()) {
        (InstanceDef::Item(_), Some(def_id)) => {
            let def = ty::WithOptConstParam::unknown(def_id);
            steal = tcx.mir_promoted(def).0;
            guard = steal.borrow();
            &*guard
        }
        // Functions from other crates and MIR shims
        _ => tcx.instance_mir(instance),
    };
    let mut calls = Vec::new();
    for bb_data in body.basic_blocks() {
        let terminator = bb_data.terminator();
        if let TerminatorKind::Call { func, .. } = &terminator.kind {
            let ty = func.ty(&body.local_decls, tcx);
            let call = match ty.kind() {
                ty::FnDef(def_id, substs) => (*def_id, *substs),
                _ => continue,
            };
            // We've seen this before
            if calls.contains(&call) {
                continue;
            }
            calls.push(call);
        }
    }
    tcx.arena.alloc_slice(&calls)
}
<commit_msg>Use small hash set in `mir_inliner_callees`<commit_after>use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use rustc_data_structures::sso::SsoHashSet;
use rustc_data_structures::stack::ensure_sufficient_stack;
use rustc_hir::def_id::{DefId, LocalDefId};
use rustc_middle::mir::TerminatorKind;
use rustc_middle::ty::TypeFoldable;
use rustc_middle::ty::{self, subst::SubstsRef, InstanceDef, TyCtxt};

// FIXME: check whether it is cheaper to precompute the entire call graph instead of invoking
// this query riddiculously often.
#[instrument(level = "debug", skip(tcx, root, target))]
crate fn mir_callgraph_reachable(
    tcx: TyCtxt<'tcx>,
    (root, target): (ty::Instance<'tcx>, LocalDefId),
) -> bool {
    trace!(%root, target = %tcx.def_path_str(target.to_def_id()));
    let param_env = tcx.param_env_reveal_all_normalized(target);
    assert_ne!(
        root.def_id().expect_local(),
        target,
        "you should not call `mir_callgraph_reachable` on immediate self recursion"
    );
    assert!(
        matches!(root.def, InstanceDef::Item(_)),
        "you should not call `mir_callgraph_reachable` on shims"
    );
    assert!(
        !tcx.is_constructor(root.def_id()),
        "you should not call `mir_callgraph_reachable` on enum/struct constructor functions"
    );
    #[instrument(
        level = "debug",
        skip(tcx, param_env, target, stack, seen, recursion_limiter, caller)
    )]
    fn process(
        tcx: TyCtxt<'tcx>,
        param_env: ty::ParamEnv<'tcx>,
        caller: ty::Instance<'tcx>,
        target: LocalDefId,
        stack: &mut Vec<ty::Instance<'tcx>>,
        seen: &mut FxHashSet<ty::Instance<'tcx>>,
        recursion_limiter: &mut FxHashMap<DefId, usize>,
    ) -> bool {
        trace!(%caller);
        for &(callee, substs) in tcx.mir_inliner_callees(caller.def) {
            let substs = caller.subst_mir_and_normalize_erasing_regions(tcx, param_env, substs);
            let callee = match ty::Instance::resolve(tcx, param_env, callee, substs).unwrap() {
                Some(callee) => callee,
                None => {
                    trace!(?callee, "cannot resolve, skipping");
                    continue;
                }
            };

            // Found a path.
            if callee.def_id() == target.to_def_id() {
                return true;
            }

            if tcx.is_constructor(callee.def_id()) {
                trace!("constructors always have MIR");
                // Constructor functions cannot cause a query cycle.
                continue;
            }

            match callee.def {
                InstanceDef::Item(_) => {
                    // If there is no MIR available (either because it was not in metadata or
                    // because it has no MIR because it's an extern function), then the inliner
                    // won't cause cycles on this.
                    if !tcx.is_mir_available(callee.def_id()) {
                        trace!(?callee, "no mir available, skipping");
                        continue;
                    }
                }
                // These have no own callable MIR.
                InstanceDef::Intrinsic(_) | InstanceDef::Virtual(..) => continue,
                // These have MIR and if that MIR is inlined, substituted and then inlining is run
                // again, a function item can end up getting inlined. Thus we'll be able to cause
                // a cycle that way
                InstanceDef::VtableShim(_)
                | InstanceDef::ReifyShim(_)
                | InstanceDef::FnPtrShim(..)
                | InstanceDef::ClosureOnceShim { .. }
                | InstanceDef::CloneShim(..) => {}
                InstanceDef::DropGlue(..) => {
                    // FIXME: A not fully substituted drop shim can cause ICEs if one attempts to
                    // have its MIR built. Likely oli-obk just screwed up the `ParamEnv`s, so this
                    // needs some more analysis.
                    if callee.needs_subst() {
                        continue;
                    }
                }
            }

            if seen.insert(callee) {
                let recursion = recursion_limiter.entry(callee.def_id()).or_default();
                trace!(?callee, recursion = *recursion);
                if tcx.sess.recursion_limit().value_within_limit(*recursion) {
                    *recursion += 1;
                    stack.push(callee);
                    let found_recursion = ensure_sufficient_stack(|| {
                        process(tcx, param_env, callee, target, stack, seen, recursion_limiter)
                    });
                    if found_recursion {
                        return true;
                    }
                    stack.pop();
                } else {
                    // Pessimistically assume that there could be recursion.
                    return true;
                }
            }
        }
        false
    }
    process(
        tcx,
        param_env,
        root,
        target,
        &mut Vec::new(),
        &mut FxHashSet::default(),
        &mut FxHashMap::default(),
    )
}

crate fn mir_inliner_callees<'tcx>(
    tcx: TyCtxt<'tcx>,
    instance: ty::InstanceDef<'tcx>,
) -> &'tcx [(DefId, SubstsRef<'tcx>)] {
    let steal;
    let guard;
    let body = match (instance, instance.def_id().as_local()) {
        (InstanceDef::Item(_), Some(def_id)) => {
            let def = ty::WithOptConstParam::unknown(def_id);
            steal = tcx.mir_promoted(def).0;
            guard = steal.borrow();
            &*guard
        }
        // Functions from other crates and MIR shims
        _ => tcx.instance_mir(instance),
    };
    let mut calls = SsoHashSet::new();
    for bb_data in body.basic_blocks() {
        let terminator = bb_data.terminator();
        if let TerminatorKind::Call { func, .. } = &terminator.kind {
            let ty = func.ty(&body.local_decls, tcx);
            let call = match ty.kind() {
                ty::FnDef(def_id, substs) => (*def_id, *substs),
                _ => continue,
            };
            calls.insert(call);
        }
    }
    tcx.arena.alloc_from_iter(calls.iter().copied())
}
<|endoftext|>