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_codegen_ssa/mir/block.rs

@@ -619,9 +619,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
@@ -642,7 +647,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_ssa/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);

src/librustc_data_structures/sync.rs

@@ -22,6 +22,7 @@ use std::collections::HashMap;
 use std::hash::{BuildHasher, Hash};
 use std::marker::PhantomData;
 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;
@@ -782,3 +783,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_middle/mir/mod.rs

@@ -49,7 +49,7 @@ mod type_foldable;
 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>;
@@ -2664,3 +2664,33 @@ impl Location {
         }
     }
 }
+
+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 BodyAndCache<'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>;
+}

src/librustc_middle/query/mod.rs

@@ -205,6 +205,21 @@ 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
+            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::BodyAndCache<'tcx>> {
+            anon
+
+            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::BodyAndCache<'tcx>> {
             cache_on_disk_if { key.is_local() }
             load_cached(tcx, id) {

src/librustc_middle/ty/mod.rs

@@ -3000,17 +3000,23 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     /// Returns the possibly-auto-generated MIR of a `(DefId, Subst)` pair.
-    pub fn instance_mir(self, instance: ty::InstanceDef<'tcx>) -> ReadOnlyBodyAndCache<'tcx, 'tcx> {
-        match instance {
+    pub fn instance_mir(self, instance: ty::Instance<'tcx>) -> ReadOnlyBodyAndCache<'tcx, 'tcx> {
+        match instance.def {
             ty::InstanceDef::Item(did) => self.optimized_mir(did).unwrap_read_only(),
+            ty::InstanceDef::Intrinsic(..) => {
+                if let Some(mir) = self.custom_intrinsic_mir(instance) {
+                    mir.unwrap_read_only()
+                } else {
+                    self.mir_shims(instance.def).unwrap_read_only()
+                }
+            }
             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).unwrap_read_only(),
+            | ty::InstanceDef::CloneShim(..) => self.mir_shims(instance.def).unwrap_read_only(),
         }
     }
 

src/librustc_mir/const_eval/eval_queries.rs

@@ -306,7 +306,7 @@ pub fn const_eval_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);
     res.and_then(|body| eval_body_using_ecx(&mut ecx, cid, &body))
         .and_then(|place| {
             Ok(RawConst { alloc_id: place.ptr.assert_ptr().alloc_id, ty: place.layout.ty })

src/librustc_mir/const_eval/fn_queries.rs

@@ -125,6 +125,9 @@ fn is_const_intrinsic(tcx: TyCtxt<'_>, def_id: DefId) -> Option<bool> {
 
     match tcx.fn_sig(def_id).abi() {
         Abi::RustIntrinsic | Abi::PlatformIntrinsic => {
+            if tcx.custom_intrinsic_mirgen(def_id).is_some() {
+                return Some(true);
+            }
             Some(tcx.lookup_const_stability(def_id).is_some())
         }
         _ => None,

src/librustc_mir/const_eval/machine.rs

@@ -218,7 +218,7 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter {
             }
         }
         // 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 {

src/librustc_mir/interpret/eval_context.rs

@@ -351,7 +351,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, mir::ReadOnlyBodyAndCache<'tcx, 'tcx>> {
         // do not continue if typeck errors occurred (can only occur in local crate)
@@ -366,7 +366,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].unwrap_read_only());
         }
-        match instance {
+        match instance.def {
             ty::InstanceDef::Item(def_id) => {
                 if self.tcx.is_mir_available(did) {
                     Ok(self.tcx.optimized_mir(did).unwrap_read_only())

src/librustc_mir/interpret/terminator.rs

@@ -203,6 +203,8 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             }
         };
 
+        let custom = self.tcx.custom_intrinsic_mir(instance);
+
         // ABI check
         {
             let callee_abi = {
@@ -281,6 +283,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

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

src/librustc_mir/monomorphize/collector.rs

@@ -693,6 +693,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(instance));
+            }
         }
         ty::InstanceDef::DropGlue(_, None) => {
             // Don't need to emit noop drop glue if we are calling directly.
@@ -1157,7 +1161,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, body: &body, output, instance }.visit_body(&body);
 }

src/librustc_mir/monomorphize/mod.rs

@@ -1,10 +1,79 @@
-use rustc_middle::traits;
+use rustc_index::vec::IndexVec;
+use rustc_middle::{traits, mir};
 use rustc_middle::ty::adjustment::CustomCoerceUnsized;
-use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::{self, Ty, TyCtxt, Instance};
+use rustc_middle::ty::query::Providers;
+use rustc_span::DUMMY_SP;
 
 pub mod collector;
 pub mod partitioning;
 
+pub fn provide(providers: &mut Providers<'_>) {
+    providers.custom_intrinsic_mirgen = |_, _| { None };
+    providers.custom_intrinsic_mir = custom_intrinsic_mir;
+}
+
+fn custom_intrinsic_mir<'tcx>(tcx: TyCtxt<'tcx>,
+                              instance: Instance<'tcx>)
+    -> Option<&'tcx mir::BodyAndCache<'tcx>>
+{
+    let mirgen = tcx.custom_intrinsic_mirgen(instance.def_id())?;
+
+    let ty = instance.monomorphic_ty(tcx);
+    let sig = ty.fn_sig(tcx);
+    let sig = tcx.normalize_erasing_late_bound_regions(
+        ty::ParamEnv::reveal_all(),
+        &sig,
+    );
+
+    // no var arg calls, so we can skip monomorphizing extra arguments.
+    assert!(!sig.c_variadic);
+
+    let source_scope_local_data = mir::ClearCrossCrate::Clear;
+    let source_scope = mir::SourceScopeData {
+        span: DUMMY_SP,
+        parent_scope: None,
+        local_data: source_scope_local_data,
+    };
+    let source_info = mir::SourceInfo {
+        span: DUMMY_SP,
+        scope: mir::OUTERMOST_SOURCE_SCOPE,
+    };
+
+    let mut source_scopes = IndexVec::new();
+    source_scopes.push(source_scope.clone());
+
+    let ret_decl = mir::LocalDecl::new_return_place(sig.output(), DUMMY_SP);
+    let mut local_decls = IndexVec::from_elem_n(ret_decl, 1);
+    for &arg in sig.inputs().iter() {
+        local_decls.push(mir::LocalDecl {
+            mutability: mir::Mutability::Mut,
+            local_info: mir::LocalInfo::Other,
+            ty: arg,
+            source_info,
+            internal: false,
+            user_ty: mir::UserTypeProjections::none(),
+            is_block_tail: None,
+        });
+    }
+
+    let gen = mir::Body::new(IndexVec::new(),
+                             source_scopes,
+                             local_decls,
+                             Default::default(),
+                             sig.inputs().len(),
+                             Vec::new(),
+                             source_scope.span,
+                             Vec::new(),
+                             None);
+    let mut gen = mir::BodyAndCache::new(gen);
+
+    mirgen.mirgen_simple_intrinsic(tcx, instance, &mut gen);
+    gen.ensure_predecessors();
+
+    Some(tcx.arena.alloc(gen))
+}
+
 pub fn custom_coerce_unsize_info<'tcx>(
     tcx: TyCtxt<'tcx>,
     source_ty: Ty<'tcx>,

src/librustc_symbol_mangling/lib.rs

@@ -183,6 +183,10 @@ fn compute_symbol_name(
             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_ty/ty.rs

@@ -294,7 +294,10 @@ fn instance_def_size_estimate<'tcx>(
 
     match instance_def {
         InstanceDef::Item(..) | InstanceDef::DropGlue(..) => {
-            let mir = tcx.instance_mir(instance_def);
+            let mir = tcx.instance_mir(ty::Instance {
+                def: instance_def,
+                substs: tcx.intern_substs(&[]),
+            });
             mir.basic_blocks().iter().map(|bb| bb.statements.len()).sum()
         }
         // Estimate the size of other compiler-generated shims to be 1.