Driver provided "intrinsics".

Allows drivers to define functions which are only expanded by the plugin
after all type checking, inference, etc etc, ie only once during mono
item collection.

The form of the interface with the plugins is slightly different than
what was proposed in rust-lang#51623.
Additionally, signature checking is added.

Author: DiamondLovesYou

compiler/rustc_codegen_ssa/src/coverageinfo/map.rs

@@ -50,7 +50,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
     }
 
     fn create(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>, is_used: bool) -> Self {
-        let coverageinfo = tcx.coverageinfo(instance.def);
+        let coverageinfo = tcx.coverageinfo(instance);
         debug!(
             "FunctionCoverage::create(instance={:?}) has coverageinfo={:?}. is_used={}",
             instance, coverageinfo, is_used

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -611,9 +611,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         let arg_count = fn_abi.args.len() + fn_abi.ret.is_indirect() as usize;
         let mut llargs = Vec::with_capacity(arg_count);
 
+        // Custom intrinsics are treated as-if they were normal functions here.
+        let is_custom_intrinsic = intrinsic.and_then(|_| instance )
+            .map(|instance| bx.tcx().custom_intrinsic_mir(instance).is_some() )
+            .unwrap_or_default();
+
         // Prepare the return value destination
         let ret_dest = if let Some((dest, _)) = *destination {
-            let is_intrinsic = intrinsic.is_some();
+            let is_intrinsic = intrinsic.is_some() && !is_custom_intrinsic;
             self.make_return_dest(&mut bx, dest, &fn_abi.ret, &mut llargs, is_intrinsic)
         } else {
             ReturnDest::Nothing
@@ -636,6 +641,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         match intrinsic {
             None | Some(sym::drop_in_place) => {}
             Some(sym::copy_nonoverlapping) => unreachable!(),
+            Some(_intrinsic) if is_custom_intrinsic => { },
             Some(intrinsic) => {
                 let dest = match ret_dest {
                     _ if fn_abi.ret.is_indirect() => llargs[0],

compiler/rustc_codegen_ssa/src/mir/coverageinfo.rs

@@ -31,7 +31,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         bx.add_coverage_counter(instance, id, code_region);
                     }
 
-                    let coverageinfo = bx.tcx().coverageinfo(instance.def);
+                    let coverageinfo = bx.tcx().coverageinfo(instance);
 
                     let fn_name = bx.get_pgo_func_name_var(instance);
                     let hash = bx.const_u64(function_source_hash);

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -137,7 +137,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     let llfn = cx.get_fn(instance);
 
-    let mir = cx.tcx().instance_mir(instance.def);
+    let mir = cx.tcx().instance_mir(instance);
 
     let fn_abi = FnAbi::of_instance(cx, instance, &[]);
     debug!("fn_abi: {:?}", fn_abi);

compiler/rustc_data_structures/src/sync.rs

@@ -21,6 +21,7 @@ use crate::owning_ref::{Erased, OwningRef};
 use std::collections::HashMap;
 use std::hash::{BuildHasher, Hash};
 use std::ops::{Deref, DerefMut};
+use std::sync::{Arc, atomic::AtomicPtr, atomic, };
 
 pub use std::sync::atomic::Ordering;
 pub use std::sync::atomic::Ordering::SeqCst;
@@ -598,3 +599,55 @@ impl<T> DerefMut for OneThread<T> {
         &mut self.inner
     }
 }
+
+/// Provides atomic mutability by replacing the value inside this `ArcCell`.
+/// Similar to the `crossbeam` structure of the same name.
+#[derive(Debug)]
+pub struct ArcCell<T>(AtomicPtr<T>);
+impl<T> ArcCell<T> {
+    pub fn new(v: Arc<T>) -> Self {
+        ArcCell(AtomicPtr::new(Arc::into_raw(v) as *mut _))
+    }
+    pub fn get(&self) -> Arc<T> {
+        let ptr = self.0.load(atomic::Ordering::Acquire);
+        let arc = unsafe { Arc::from_raw(ptr as *const T) };
+        let ret = arc.clone();
+        // don't drop our copy:
+        ::std::mem::forget(arc);
+        ret
+    }
+    /// Update the value, returning the previous value.
+    pub fn set(&self, v: Arc<T>) -> Arc<T> {
+        let new = Arc::into_raw(v) as *mut _;
+        let mut expected = self.0.load(atomic::Ordering::Acquire);
+        loop {
+            match self.0.compare_exchange_weak(expected, new,
+                                               atomic::Ordering::SeqCst,
+                                               atomic::Ordering::Acquire) {
+                Ok(old) => {
+                    return unsafe { Arc::from_raw(old as *const T) };
+                },
+                Err(v) => {
+                    expected = v;
+                },
+            }
+        }
+    }
+}
+impl<T> Drop for ArcCell<T> {
+    fn drop(&mut self) {
+        let ptr = self.0.load(atomic::Ordering::Acquire);
+        // drop our copy of the arc:
+        unsafe { Arc::from_raw(ptr as *const _) };
+    }
+}
+impl<T> Clone for ArcCell<T> {
+    fn clone(&self) -> Self {
+        ArcCell::new(self.get())
+    }
+}
+impl<T> From<Arc<T>> for ArcCell<T> {
+    fn from(v: Arc<T>) -> Self {
+        Self::new(v)
+    }
+}

compiler/rustc_middle/src/dep_graph/dep_node.rs

@@ -285,10 +285,10 @@ pub type DepNode = rustc_query_system::dep_graph::DepNode<DepKind>;
 // required that their size stay the same, but we don't want to change
 // it inadvertently. This assert just ensures we're aware of any change.
 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-static_assert_size!(DepNode, 17);
+static_assert_size!(DepNode, 18);
 
 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
-static_assert_size!(DepNode, 24);
+static_assert_size!(DepNode, 25);
 
 pub trait DepNodeExt: Sized {
     /// Construct a DepNode from the given DepKind and DefPathHash. This

compiler/rustc_middle/src/mir/mod.rs

@@ -2900,3 +2900,33 @@ impl Location {
         }
     }
 }
+
+pub trait CustomIntrinsicMirGen: Sync + Send + fmt::Debug {
+    /// Codegen a plugin-defined intrinsic. This is intended to be used to
+    /// "return" values based on the monomorphized and erased types of the
+    /// function call. Codegen will codegen the `extra_stmts` and then insert
+    /// an unconditional branch to the exit block.
+    ///
+    /// Consider this to be highly unstable; it will likely change without
+    /// warning. There is also no spec for this, it is 100% implementation
+    /// defined, and may not be implemented at all for some codegen backends.
+    ///
+    /// If the codegen backend is multithreaded, this will be called from
+    /// any number of threads, hence `Sync + Send`.
+    ///
+    /// YOU ARE RESPONSIBLE FOR THE SAFETY OF THE EXTRA STATEMENTS.
+    /// You have been warned. Good luck, have fun.
+    fn mirgen_simple_intrinsic<'tcx>(&self,
+                                     tcx: TyCtxt<'tcx>,
+                                     instance: ty::Instance<'tcx>,
+                                     mir: &mut Body<'tcx>);
+
+    /// The following are used for typeck-ing:
+
+    /// The number of generic parameters expected.
+    fn generic_parameter_count<'tcx>(&self, tcx: TyCtxt<'tcx>) -> usize;
+    /// The types of the input args.
+    fn inputs<'tcx>(&self, tcx: TyCtxt<'tcx>) -> &'tcx List<Ty<'tcx>>;
+    /// The return type.
+    fn output<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Ty<'tcx>;
+}

compiler/rustc_middle/src/query/mod.rs

@@ -338,7 +338,7 @@ rustc_queries! {
 
     /// Returns coverage summary info for a function, after executing the `InstrumentCoverage`
     /// MIR pass (assuming the -Zinstrument-coverage option is enabled).
-    query coverageinfo(key: ty::InstanceDef<'tcx>) -> mir::CoverageInfo {
+    query coverageinfo(key: ty::Instance<'tcx>) -> mir::CoverageInfo {
         desc { |tcx| "retrieving coverage info from MIR for `{}`", tcx.def_path_str(key.def_id()) }
         storage(ArenaCacheSelector<'tcx>)
     }
@@ -363,6 +363,20 @@ rustc_queries! {
         storage(ArenaCacheSelector<'tcx>)
         cache_on_disk_if { key.is_local() }
     }
+    /// If defined by the driver, returns the extra mir statements to codegen,
+    /// else returns `None`.
+    query custom_intrinsic_mirgen(key: DefId) -> Option<Lrc<dyn mir::CustomIntrinsicMirGen>> {
+        anon
+        no_hash
+
+        desc { |tcx| "asking for the custom MIR generator of `{}`", tcx.def_path_str(key) }
+    }
+    /// The monomorphized MIR for a custom intrinsic instance.
+    query custom_intrinsic_mir(inst: ty::Instance<'tcx>) -> Option<&'tcx mir::Body<'tcx>> {
+        anon
+
+        desc { |tcx| "asking for the custom MIR of `{}`", tcx.def_path_str(inst.def_id()) }
+    }
 
     /// The `DefId` is the `DefId` of the containing MIR body. Promoteds do not have their own
     /// `DefId`. This function returns all promoteds in the specified body. The body references
@@ -782,7 +796,7 @@ rustc_queries! {
     }
 
     /// Obtain all the calls into other local functions
-    query mir_inliner_callees(key: ty::InstanceDef<'tcx>) -> &'tcx [(DefId, SubstsRef<'tcx>)] {
+    query mir_inliner_callees(key: ty::Instance<'tcx>) -> &'tcx [(DefId, SubstsRef<'tcx>)] {
         fatal_cycle
         desc { |tcx|
             "computing all local function calls in `{}`",

compiler/rustc_middle/src/ty/mod.rs

@@ -1779,8 +1779,8 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     /// Returns the possibly-auto-generated MIR of a `(DefId, Subst)` pair.
-    pub fn instance_mir(self, instance: ty::InstanceDef<'tcx>) -> &'tcx Body<'tcx> {
-        match instance {
+    pub fn instance_mir(self, instance: ty::Instance<'tcx>) -> &'tcx Body<'tcx> {
+        match instance.def {
             ty::InstanceDef::Item(def) => match self.def_kind(def.did) {
                 DefKind::Const
                 | DefKind::Static
@@ -1794,14 +1794,20 @@ impl<'tcx> TyCtxt<'tcx> {
                     self.optimized_mir(def.did)
                 }
             },
+            ty::InstanceDef::Intrinsic(..) => {
+                if let Some(mir) = self.custom_intrinsic_mir(instance) {
+                    mir
+                } else {
+                    self.mir_shims(instance.def)
+                }
+            },
             ty::InstanceDef::VtableShim(..)
             | ty::InstanceDef::ReifyShim(..)
-            | ty::InstanceDef::Intrinsic(..)
             | ty::InstanceDef::FnPtrShim(..)
             | ty::InstanceDef::Virtual(..)
             | ty::InstanceDef::ClosureOnceShim { .. }
             | ty::InstanceDef::DropGlue(..)
-            | ty::InstanceDef::CloneShim(..) => self.mir_shims(instance),
+            | ty::InstanceDef::CloneShim(..) => self.mir_shims(instance.def),
         }
     }
 

compiler/rustc_mir/src/const_eval/eval_queries.rs

@@ -306,7 +306,7 @@ pub fn eval_to_allocation_raw_provider<'tcx>(
         MemoryExtra { can_access_statics: is_static },
     );
 
-    let res = ecx.load_mir(cid.instance.def, cid.promoted);
+    let res = ecx.load_mir(cid.instance, cid.promoted);
     match res.and_then(|body| eval_body_using_ecx(&mut ecx, cid, &body)) {
         Err(error) => {
             let err = ConstEvalErr::new(&ecx, error, None);

compiler/rustc_mir/src/const_eval/machine.rs

@@ -203,9 +203,9 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
 
     fn load_mir(
         ecx: &InterpCx<'mir, 'tcx, Self>,
-        instance: ty::InstanceDef<'tcx>,
+        instance: ty::Instance<'tcx>,
     ) -> InterpResult<'tcx, &'tcx mir::Body<'tcx>> {
-        match instance {
+        match instance.def {
             ty::InstanceDef::Item(def) => {
                 if ecx.tcx.is_ctfe_mir_available(def.did) {
                     Ok(ecx.tcx.mir_for_ctfe_opt_const_arg(def))
@@ -242,7 +242,7 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
             }
         }
         // This is a const fn. Call it.
-        Ok(Some(match ecx.load_mir(instance.def, None) {
+        Ok(Some(match ecx.load_mir(instance, None) {
             Ok(body) => body,
             Err(err) => {
                 if let err_unsup!(NoMirFor(did)) = err.kind() {

compiler/rustc_mir/src/interpret/eval_context.rs

@@ -485,11 +485,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
 
     pub fn load_mir(
         &self,
-        instance: ty::InstanceDef<'tcx>,
+        instance: ty::Instance<'tcx>,
         promoted: Option<mir::Promoted>,
     ) -> InterpResult<'tcx, &'tcx mir::Body<'tcx>> {
         // do not continue if typeck errors occurred (can only occur in local crate)
-        let def = instance.with_opt_param();
+        let def = instance.def.with_opt_param();
         if let Some(def) = def.as_local() {
             if self.tcx.has_typeck_results(def.did) {
                 if let Some(error_reported) = self.tcx.typeck_opt_const_arg(def).tainted_by_errors {

compiler/rustc_mir/src/interpret/machine.rs

@@ -142,7 +142,7 @@ pub trait Machine<'mir, 'tcx>: Sized {
     /// constants, ...
     fn load_mir(
         ecx: &InterpCx<'mir, 'tcx, Self>,
-        instance: ty::InstanceDef<'tcx>,
+        instance: ty::Instance<'tcx>,
     ) -> InterpResult<'tcx, &'tcx mir::Body<'tcx>> {
         Ok(ecx.tcx.instance_mir(instance))
     }

compiler/rustc_mir/src/interpret/terminator.rs

@@ -256,6 +256,8 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             }
         };
 
+        let custom = self.tcx.custom_intrinsic_mir(instance);
+
         let get_abi = |this: &Self, instance_ty: Ty<'tcx>| match instance_ty.kind() {
             ty::FnDef(..) => instance_ty.fn_sig(*this.tcx).abi(),
             ty::Closure(..) => Abi::RustCall,
@@ -355,6 +357,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                     // The Rust ABI is special: ZST get skipped.
                     let rust_abi = match caller_abi {
                         Abi::Rust | Abi::RustCall => true,
+                        Abi::RustIntrinsic if custom.is_some() => true,
                         _ => false,
                     };
                     // We have two iterators: Where the arguments come from,

compiler/rustc_mir/src/lib.rs

@@ -52,6 +52,7 @@ pub fn provide(providers: &mut Providers) {
     const_eval::provide(providers);
     shim::provide(providers);
     transform::provide(providers);
+    monomorphize::provide(providers);
     monomorphize::partitioning::provide(providers);
     monomorphize::polymorphize::provide(providers);
     providers.eval_to_const_value_raw = const_eval::eval_to_const_value_raw_provider;

compiler/rustc_mir/src/monomorphize/collector.rs

@@ -907,6 +907,10 @@ fn visit_instance_use<'tcx>(
             if !is_direct_call {
                 bug!("{:?} being reified", instance);
             }
+
+            if let Some(_mir) = tcx.custom_intrinsic_mir(instance) {
+                output.push(create_fn_mono_item(tcx, instance, source));
+            }
         }
         ty::InstanceDef::DropGlue(_, None) => {
             // Don't need to emit noop drop glue if we are calling directly.
@@ -1374,7 +1378,7 @@ fn collect_neighbours<'tcx>(
     output: &mut Vec<Spanned<MonoItem<'tcx>>>,
 ) {
     debug!("collect_neighbours: {:?}", instance.def_id());
-    let body = tcx.instance_mir(instance.def);
+    let body = tcx.instance_mir(instance);
 
     MirNeighborCollector { tcx, body: &body, output, instance }.visit_body(&body);
 }