Ensure exclusive access to diagnostics

crates/rune/src/no_std/mod.rs

@@ -1,22 +1,22 @@
 //! Public types related to using rune in #[no_std] environments.
 
-use crate::runtime::vm_diagnostics::VmDiagnosticsObj;
+use core::ptr::NonNull;
 
 /// Environment that needs to be stored somewhere.
 #[derive(Clone, Copy)]
 #[repr(C)]
 pub struct RawEnv {
-    pub(crate) context: *const (),
-    pub(crate) unit: *const (),
-    pub(crate) diagnostics: Option<VmDiagnosticsObj>,
+    pub(crate) context: Option<NonNull<()>>,
+    pub(crate) unit: Option<NonNull<()>>,
+    pub(crate) diagnostics: Option<NonNull<()>>,
 }
 
 impl RawEnv {
     /// Initialize an empty raw environment.
     pub const fn null() -> RawEnv {
         RawEnv {
-            context: core::ptr::null(),
-            unit: core::ptr::null(),
+            context: None,
+            unit: None,
             diagnostics: None,
         }
     }

crates/rune/src/runtime/env.rs

@@ -7,6 +7,8 @@
 //!
 //! See the corresponding function for documentation.
 
+use core::ptr::NonNull;
+
 #[cfg_attr(feature = "std", path = "env/std.rs")]
 mod no_std;
 
@@ -15,34 +17,70 @@ use ::rust_alloc::sync::Arc;
 use crate::runtime::vm_diagnostics::VmDiagnosticsObj;
 use crate::runtime::{RuntimeContext, Unit, VmErrorKind, VmResult};
 
-/// Call the given closure with access to the checked environment.
-pub(crate) fn with<F, T>(c: F) -> VmResult<T>
+/// Access shared parts of the environment.
+///
+/// This does not take ownership of the environment, so the environment can be
+/// recursively accessed.
+pub(crate) fn shared<F, T>(c: F) -> VmResult<T>
 where
-    F: FnOnce(&Arc<RuntimeContext>, &Arc<Unit>, Option<&mut VmDiagnosticsObj>) -> VmResult<T>,
+    F: FnOnce(&Arc<RuntimeContext>, &Arc<Unit>) -> VmResult<T>,
 {
     let env = self::no_std::rune_env_get();
+
     let Env {
-        context,
-        unit,
-        mut diagnostics,
-    } = env;
+        context: Some(context),
+        unit: Some(unit),
+        ..
+    } = env
+    else {
+        return VmResult::err(VmErrorKind::MissingInterfaceEnvironment);
+    };
+
+    // Safety: context and unit can only be registered publicly through
+    // [`Guard`], which makes sure that they are live for the duration of the
+    // registration.
+    let context = unsafe { context.as_ref() };
+    let unit = unsafe { unit.as_ref() };
+    c(context, unit)
+}
 
-    if context.is_null() || unit.is_null() {
+/// Call the given closure with access to the checked environment accessing it
+/// exclusively.
+///
+/// This takes ownership of the environment, so recursive calls are not
+/// supported.
+pub(crate) fn exclusive<F, T>(c: F) -> VmResult<T>
+where
+    F: FnOnce(&Arc<RuntimeContext>, &Arc<Unit>, Option<&mut VmDiagnosticsObj>) -> VmResult<T>,
+{
+    let guard = Guard {
+        env: self::no_std::rune_env_replace(Env::null()),
+    };
+
+    let Env {
+        context: Some(context),
+        unit: Some(unit),
+        ..
+    } = guard.env
+    else {
         return VmResult::err(VmErrorKind::MissingInterfaceEnvironment);
-    }
+    };
 
     // Safety: context and unit can only be registered publicly through
-    // [Guard], which makes sure that they are live for the duration of
-    // the registration.
-    c(
-        unsafe { &*context },
-        unsafe { &*unit },
-        diagnostics.as_mut(),
-    )
+    // [`Guard`], which makes sure that they are live for the duration of the
+    // registration.
+    let context = unsafe { context.as_ref() };
+    let unit = unsafe { unit.as_ref() };
+    let diagnostics = match guard.env.diagnostics {
+        Some(mut d) => Some(unsafe { d.as_mut() }),
+        None => None,
+    };
+
+    c(context, unit, diagnostics)
 }
 
 pub(crate) struct Guard {
-    old: Env,
+    env: Env,
 }
 
 impl Guard {
@@ -52,38 +90,37 @@ impl Guard {
     ///
     /// The returned guard must be dropped before the pointed to elements are.
     pub(crate) fn new(
-        context: *const Arc<RuntimeContext>,
-        unit: *const Arc<Unit>,
-        diagnostics: Option<VmDiagnosticsObj>,
+        context: NonNull<Arc<RuntimeContext>>,
+        unit: NonNull<Arc<Unit>>,
+        diagnostics: Option<NonNull<VmDiagnosticsObj>>,
     ) -> Guard {
-        let old = self::no_std::rune_env_replace(Env {
-            context,
-            unit,
+        let env = self::no_std::rune_env_replace(Env {
+            context: Some(context),
+            unit: Some(unit),
             diagnostics,
         });
-        Guard { old }
+        Guard { env }
     }
 }
 
 impl Drop for Guard {
     fn drop(&mut self) {
-        let _ = self::no_std::rune_env_replace(self.old);
+        let _ = self::no_std::rune_env_replace(self.env);
     }
 }
 
 #[derive(Debug, Clone, Copy)]
 struct Env {
-    context: *const Arc<RuntimeContext>,
-    unit: *const Arc<Unit>,
-    diagnostics: Option<VmDiagnosticsObj>,
+    context: Option<NonNull<Arc<RuntimeContext>>>,
+    unit: Option<NonNull<Arc<Unit>>>,
+    diagnostics: Option<NonNull<VmDiagnosticsObj>>,
 }
 
 impl Env {
-    #[cfg(feature = "std")]
     const fn null() -> Self {
         Self {
-            context: core::ptr::null(),
-            unit: core::ptr::null(),
+            context: None,
+            unit: None,
             diagnostics: None,
         }
     }

crates/rune/src/runtime/env/no_std.rs

@@ -30,16 +30,16 @@ pub(super) fn rune_env_replace(env: Env) -> Env {
 
 unsafe fn from_env(env: Env) -> RawEnv {
     RawEnv {
-        context: env.context as *const _,
-        unit: env.unit as *const _,
-        diagnostics: env.diagnostics,
+        context: env.context.map(|ptr| ptr.cast()),
+        unit: env.unit.map(|ptr| ptr.cast()),
+        diagnostics: env.diagnostics.map(|ptr| ptr.cast()),
     }
 }
 
 unsafe fn from_raw_env(env: RawEnv) -> Env {
     Env {
-        context: env.context as *const _,
-        unit: env.unit as *const _,
-        diagnostics: env.diagnostics,
+        context: env.context.map(|ptr| ptr.cast()),
+        unit: env.unit.map(|ptr| ptr.cast()),
+        diagnostics: env.diagnostics.map(|ptr| ptr.cast()),
     }
 }

crates/rune/src/runtime/protocol_caller.rs

@@ -43,7 +43,7 @@ impl ProtocolCaller for EnvProtocolCaller {
     where
         A: GuardedArgs,
     {
-        return crate::runtime::env::with(|context, unit, _| {
+        return crate::runtime::env::shared(|context, unit| {
             let count = args.count() + 1;
             let hash = Hash::associated_function(vm_try!(target.type_hash()), protocol.hash);
 

crates/rune/src/runtime/value.rs

@@ -898,7 +898,7 @@ impl Value {
     pub fn into_type_name(self) -> VmResult<String> {
         let hash = Hash::associated_function(vm_try!(self.type_hash()), Protocol::INTO_TYPE_NAME);
 
-        crate::runtime::env::with(|context, unit, _| {
+        crate::runtime::env::shared(|context, unit| {
             if let Some(name) = context.constant(hash) {
                 match name {
                     ConstValue::String(s) => {

crates/rune/src/runtime/vm.rs

@@ -1,6 +1,7 @@
 use core::cmp::Ordering;
 use core::mem::{replace, swap};
 use core::ops;
+use core::ptr::NonNull;
 use core::slice;
 
 use ::rust_alloc::sync::Arc;
@@ -613,7 +614,7 @@ impl Vm {
     }
 
     fn called_function_hook(&self, hash: Hash) -> VmResult<()> {
-        crate::runtime::env::with(|_, _, diagnostics| {
+        crate::runtime::env::exclusive(|_, _, diagnostics| {
             if let Some(diagnostics) = diagnostics {
                 vm_try!(diagnostics.function_used(hash, self.ip()));
             }
@@ -3065,17 +3066,33 @@ impl Vm {
     where
         F: FnOnce() -> T,
     {
-        let _guard = crate::runtime::env::Guard::new(&self.context, &self.unit, None);
+        let _guard = crate::runtime::env::Guard::new(
+            NonNull::from(&self.context),
+            NonNull::from(&self.unit),
+            None,
+        );
         f()
     }
 
     /// Evaluate a single instruction.
     pub(crate) fn run(&mut self, diagnostics: Option<&mut dyn VmDiagnostics>) -> VmResult<VmHalt> {
-        let diagnostics = diagnostics.map(VmDiagnosticsObj::new);
+        let mut vm_diagnostics_obj;
+
+        let diagnostics = match diagnostics {
+            Some(diagnostics) => {
+                vm_diagnostics_obj = VmDiagnosticsObj::new(diagnostics);
+                Some(NonNull::from(&mut vm_diagnostics_obj))
+            }
+            None => None,
+        };
 
         // NB: set up environment so that native function can access context and
         // unit.
-        let _guard = crate::runtime::env::Guard::new(&self.context, &self.unit, diagnostics);
+        let _guard = crate::runtime::env::Guard::new(
+            NonNull::from(&self.context),
+            NonNull::from(&self.unit),
+            diagnostics,
+        );
 
         loop {
             if !budget::take() {