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

src/librustc/mir/mod.rs

@@ -49,7 +49,7 @@ pub mod traversal;
 pub mod visit;
 
 /// Types for locals
-type LocalDecls<'tcx> = IndexVec<Local, LocalDecl<'tcx>>;
+pub type LocalDecls<'tcx> = IndexVec<Local, LocalDecl<'tcx>>;
 
 pub trait HasLocalDecls<'tcx> {
     fn local_decls(&self) -> &LocalDecls<'tcx>;
@@ -3350,3 +3350,41 @@ impl<'tcx> TypeFoldable<'tcx> for Constant<'tcx> {
         self.literal.visit_with(visitor)
     }
 }
+
+pub trait CustomIntrinsicMirGen: Sync + Send {
+    /// 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>;
+}
+
+/*impl<'a> HashStable<StableHashingContext<'a>> for dyn CustomIntrinsicMirGen {
+    fn hash_stable<W: StableHasherResult>(&self,
+                                          _ctx: &mut StableHashingContext<'_>,
+                                          _hasher: &mut StableHasher<W>) {
+        // TO DO
+    }
+}*/

src/librustc/query/mod.rs

@@ -129,6 +129,24 @@ rustc_queries! {
             }
         }
 
+        /// 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
+            fatal_cycle
+            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
+            fatal_cycle
+            no_force
+
+            desc { |tcx| "asking for the custom MIR of `{}`", tcx.def_path_str(inst.def_id()) }
+        }
+
         query promoted_mir(key: DefId) -> &'tcx IndexVec<mir::Promoted, mir::Body<'tcx>> {
             cache_on_disk_if { key.is_local() }
             load_cached(tcx, id) {

src/librustc/ty/mod.rs

@@ -3032,20 +3032,26 @@ 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(did) => {
                 self.optimized_mir(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)
+                self.mir_shims(instance.def)
             }
         }
     }
@@ -3329,7 +3335,12 @@ fn instance_def_size_estimate<'tcx>(tcx: TyCtxt<'tcx>, instance_def: InstanceDef
     match instance_def {
         InstanceDef::Item(..) |
         InstanceDef::DropGlue(..) => {
-            let mir = tcx.instance_mir(instance_def);
+            let instance = Instance {
+                def: instance_def,
+                // this field can be whatever because it won't be used in this case.
+                substs: tcx.intern_substs(&[]),
+            };
+            let mir = tcx.instance_mir(instance);
             mir.basic_blocks().iter().map(|bb| bb.statements.len()).sum()
         },
         // Estimate the size of other compiler-generated shims to be 1.

src/librustc_codegen_ssa/base.rs

@@ -370,7 +370,7 @@ pub fn codegen_instance<'a, 'tcx: 'a, Bx: BuilderMethods<'a, 'tcx>>(
 
     let lldecl = cx.get_fn(instance);
 
-    let mir = cx.tcx().instance_mir(instance.def);
+    let mir = cx.tcx().instance_mir(instance);
     mir::codegen_mir::<Bx>(cx, lldecl, &mir, instance, sig);
 }
 

src/librustc_codegen_ssa/mir/block.rs

@@ -570,9 +570,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((ref 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 {
@@ -594,7 +599,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             return;
         }
 
-        if intrinsic.is_some() && intrinsic != Some("drop_in_place") {
+        if intrinsic.is_some() && intrinsic != Some("drop_in_place") &&
+            !is_custom_intrinsic {
             let dest = match ret_dest {
                 _ if fn_abi.ret.is_indirect() => llargs[0],
                 ReturnDest::Nothing =>

src/librustc_codegen_utils/symbol_names.rs

@@ -137,6 +137,10 @@ fn symbol_name(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Symbol {
             Node::ForeignItem(_) => true,
             _ => false,
         }
+    } else if let ty::InstanceDef::Intrinsic(_) = instance.def {
+        // custom intrinsics should never be foreign, otherwise
+        // generic parameters will cause duplicate symbols names.
+        tcx.custom_intrinsic_mir(instance).is_none()
     } else {
         tcx.is_foreign_item(def_id)
     };

src/librustc_data_structures/sync.rs

@@ -21,6 +21,7 @@ use std::collections::HashMap;
 use std::hash::{Hash, BuildHasher};
 use std::marker::PhantomData;
 use std::ops::{Deref, DerefMut};
+use std::sync::{Arc, atomic::AtomicPtr, atomic, };
 use crate::owning_ref::{Erased, OwningRef};
 
 pub use std::sync::atomic::Ordering::SeqCst;
@@ -790,3 +791,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)
+    }
+}

src/librustc_mir/const_eval.rs

@@ -327,8 +327,13 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
         ret: Option<mir::BasicBlock>,
     ) -> InterpResult<'tcx, Option<&'mir mir::Body<'tcx>>> {
         debug!("eval_fn_call: {:?}", instance);
+        // WIP: We assume all custom intrinsics are const. This is okay for now
+        // as these are generated in the provider/query system, which caches the
+        // resulting MIR. Effectively, this means for any specific set of generic
+        // params, the MIR is the same for each call (ie in user code).
+        let custom = ecx.tcx.custom_intrinsic_mir(instance);
         // Only check non-glue functions
-        if let ty::InstanceDef::Item(def_id) = instance.def {
+        if let (None, ty::InstanceDef::Item(def_id)) = (custom, instance.def) {
             // Execution might have wandered off into other crates, so we cannot do a stability-
             // sensitive check here.  But we can at least rule out functions that are not const
             // at all.
@@ -345,7 +350,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(ref path)) = err.kind {
@@ -692,7 +697,7 @@ pub fn const_eval_raw_provider<'tcx>(
         Default::default()
     );
 
-    let res = ecx.load_mir(cid.instance.def, cid.promoted);
+    let res = ecx.load_mir(cid.instance, cid.promoted);
     res.and_then(
         |body| eval_body_using_ecx(&mut ecx, cid, body)
     ).and_then(|place| {

src/librustc_mir/interpret/eval_context.rs

@@ -282,7 +282,7 @@ impl<'mir, 'tcx, 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)
@@ -297,7 +297,7 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         if let Some(promoted) = promoted {
             return Ok(&self.tcx.promoted_mir(did)[promoted]);
         }
-        match instance {
+        match instance.def {
             ty::InstanceDef::Item(def_id) => if self.tcx.is_mir_available(did) {
                 Ok(self.tcx.optimized_mir(did))
             } else {

src/librustc_mir/interpret/terminator.rs

@@ -247,8 +247,10 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             }
         };
 
+        let custom = self.tcx.custom_intrinsic_mir(instance);
+
         match instance.def {
-            ty::InstanceDef::Intrinsic(..) => {
+            ty::InstanceDef::Intrinsic(..) if custom.is_none() => {
                 // The intrinsic itself cannot diverge, so if we got here without a return
                 // place... (can happen e.g., for transmute returning `!`)
                 let dest = match dest {
@@ -262,6 +264,7 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 self.dump_place(*dest);
                 Ok(())
             }
+            ty::InstanceDef::Intrinsic(..) |
             ty::InstanceDef::VtableShim(..) |
             ty::InstanceDef::ReifyShim(..) |
             ty::InstanceDef::ClosureOnceShim { .. } |
@@ -332,6 +335,7 @@ impl<'mir, 'tcx, 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,

src/librustc_mir/lib.rs

@@ -54,6 +54,7 @@ pub fn provide(providers: &mut Providers<'_>) {
     borrow_check::provide(providers);
     shim::provide(providers);
     transform::provide(providers);
+    monomorphize::provide(providers);
     monomorphize::partitioning::provide(providers);
     providers.const_eval = const_eval::const_eval_provider;
     providers.const_eval_raw = const_eval::const_eval_raw_provider;

src/librustc_mir/monomorphize/collector.rs

@@ -747,6 +747,10 @@ fn visit_instance_use<'tcx>(
             if !is_direct_call {
                 bug!("intrinsic {:?} being reified", def_id);
             }
+
+            if let Some(_mir) = tcx.custom_intrinsic_mir(instance) {
+                output.push(create_fn_mono_item(instance));
+            }
         }
         ty::InstanceDef::VtableShim(..) |
         ty::InstanceDef::ReifyShim(..) |
@@ -1251,7 +1255,7 @@ fn collect_neighbours<'tcx>(
     output: &mut Vec<MonoItem<'tcx>>,
 ) {
     debug!("collect_neighbours: {:?}", instance.def_id());
-    let body = tcx.instance_mir(instance.def);
+    let body = tcx.instance_mir(instance);
 
     MirNeighborCollector {
         tcx,