Add global and generational roots

This requires ocaml-sys 0.19.1, which exposes the required runtime API.
tizoc · Mar 5, 2021 · 296fe29 · 296fe29
1 parent d2c1713
commit 296fe29
Show file tree

Hide file tree

Showing 6 changed files with 176 additions and 9 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,7 +16,7 @@ exclude = [
 features = [ "without-ocamlopt" ]
 
 [dependencies]
-ocaml-sys = "^0.19"
+ocaml-sys = "^0.19.1"
 static_assertions = "1.1.0"
 
 [features]

diff --git a/src/lib.rs b/src/lib.rs
@@ -356,7 +356,7 @@ mod value;
 pub use crate::closure::{OCamlFn1, OCamlFn2, OCamlFn3, OCamlFn4, OCamlFn5};
 pub use crate::conv::{FromOCaml, ToOCaml};
 pub use crate::error::OCamlException;
-pub use crate::memory::OCamlRef;
+pub use crate::memory::{OCamlGenerationalRoot, OCamlGlobalRoot, OCamlRef};
 pub use crate::mlvalues::{
     OCamlBytes, OCamlFloat, OCamlInt, OCamlInt32, OCamlInt64, OCamlList, RawOCaml,
 };

diff --git a/src/memory.rs b/src/memory.rs
@@ -7,7 +7,12 @@ use crate::{
     runtime::OCamlRuntime,
     value::OCaml,
 };
-use core::{cell::UnsafeCell, marker::PhantomData, ptr};
+use core::{
+    cell::{Cell, UnsafeCell},
+    marker::PhantomData,
+    pin::Pin,
+    ptr,
+};
 pub use ocaml_sys::{
     caml_alloc, local_roots as ocaml_sys_local_roots, set_local_roots as ocaml_sys_set_local_roots,
     store_field,
@@ -122,6 +127,105 @@ impl<'a> OCamlRawRoot<'a> {
     }
 }
 
+/// A global root for keeping OCaml values alive and tracked
+///
+/// This allows keeping a value around when exiting the stack frame.
+///
+/// See [`OCaml::register_global_root`].
+pub struct OCamlGlobalRoot<T> {
+    pub(crate) cell: Pin<Box<Cell<RawOCaml>>>,
+    _marker: PhantomData<Cell<T>>,
+}
+
+impl<T> std::fmt::Debug for OCamlGlobalRoot<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "OCamlGlobalRoot({:#x})", self.cell.get())
+    }
+}
+
+impl<T> OCamlGlobalRoot<T> {
+    // NOTE: we require initialisation here, unlike OCamlRoot which delays it
+    // This is because we register with the GC in the constructor,
+    // for easy pairing with Drop, and registering without initializing
+    // would break OCaml runtime invariants.
+    // Always registering with UNIT (like for GCFrame initialisation)
+    // would also work, but for OCamlGenerationalRoot that would
+    // make things slower (updating requires notifying the GC),
+    // and it's better if the API is the same for both kinds of global roots.
+    pub(crate) fn new(val: OCaml<T>) -> Self {
+        let r = Self {
+            cell: Box::pin(Cell::new(val.raw)),
+            _marker: PhantomData,
+        };
+        unsafe { ocaml_sys::caml_register_global_root(r.cell.as_ptr()) };
+        r
+    }
+
+    /// Access the rooted value
+    pub fn get_ref(&self) -> OCamlRef<T> {
+        unsafe { OCamlCell::create_ref(self.cell.as_ptr()) }
+    }
+
+    /// Replace the rooted value
+    pub fn set(&self, val: OCaml<T>) {
+        self.cell.replace(val.raw);
+    }
+}
+
+impl<T> Drop for OCamlGlobalRoot<T> {
+    fn drop(&mut self) {
+        unsafe { ocaml_sys::caml_remove_global_root(self.cell.as_ptr()) };
+    }
+}
+
+/// A global, GC-friendly root for keeping OCaml values alive and tracked
+///
+/// This allows keeping a value around when exiting the stack frame.
+///
+/// Unlike with [`OCamlGlobalRoot`], the GC doesn't have to walk
+/// referenced values on every minor collection.  This makes collection
+/// faster, except if the value is short-lived and frequently updated.
+///
+/// See [`OCaml::register_generational_root`].
+pub struct OCamlGenerationalRoot<T> {
+    pub(crate) cell: Pin<Box<Cell<RawOCaml>>>,
+    _marker: PhantomData<Cell<T>>,
+}
+
+impl<T> std::fmt::Debug for OCamlGenerationalRoot<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "OCamlGenerationalRoot({:#x})", self.cell.get())
+    }
+}
+
+impl<T> OCamlGenerationalRoot<T> {
+    pub(crate) fn new(val: OCaml<T>) -> Self {
+        let r = Self {
+            cell: Box::pin(Cell::new(val.raw)),
+            _marker: PhantomData,
+        };
+        unsafe { ocaml_sys::caml_register_generational_global_root(r.cell.as_ptr()) };
+        r
+    }
+
+    /// Access the rooted value
+    pub fn get_ref(&self) -> OCamlRef<T> {
+        unsafe { OCamlCell::create_ref(self.cell.as_ptr()) }
+    }
+
+    /// Replace the rooted value
+    pub fn set(&self, val: OCaml<T>) {
+        unsafe { ocaml_sys::caml_modify_generational_global_root(self.cell.as_ptr(), val.raw) };
+        debug_assert_eq!(self.cell.get(), val.raw);
+    }
+}
+
+impl<T> Drop for OCamlGenerationalRoot<T> {
+    fn drop(&mut self) {
+        unsafe { ocaml_sys::caml_remove_generational_global_root(self.cell.as_ptr()) };
+    }
+}
+
 pub struct OCamlCell<T> {
     cell: UnsafeCell<RawOCaml>,
     _marker: PhantomData<T>,

diff --git a/src/value.rs b/src/value.rs
@@ -1,6 +1,7 @@
 // Copyright (c) SimpleStaking and Tezedge Contributors
 // SPDX-License-Identifier: MIT
 
+use crate::memory::{OCamlGenerationalRoot, OCamlGlobalRoot};
 use crate::{
     error::OCamlFixnumConversionError, memory::OCamlCell, mlvalues::*, FromOCaml, OCamlRef,
     OCamlRuntime,
@@ -109,6 +110,20 @@ impl<'a, T> OCaml<'a, T> {
     {
         RustT::from_ocaml(*self)
     }
+
+    /// Register a global root with the OCaml runtime
+    ///
+    /// If the value is seldom modified ([`OCamlGlobalRoot::set`] isn't
+    /// frequently used), [`OCaml::register_generational_root`] can be
+    /// faster.
+    pub fn register_global_root(self) -> OCamlGlobalRoot<T> {
+        OCamlGlobalRoot::new(self)
+    }
+
+    /// Register a GC-friendly global root with the OCaml runtime
+    pub fn register_generational_root(self) -> OCamlGenerationalRoot<T> {
+        OCamlGenerationalRoot::new(self)
+    }
 }
 
 impl OCaml<'static, ()> {

diff --git a/testing/rust-caller/build.rs b/testing/rust-caller/build.rs
@@ -8,7 +8,13 @@ fn main() {
     let ocaml_callable_dir = "./ocaml";
     let dune_dir = "../../_build/default/testing/rust-caller/ocaml";
     Command::new("opam")
-        .args(&["exec", "--", "dune", "build", &format!("{}/callable.exe.o", ocaml_callable_dir)])
+        .args(&[
+            "exec",
+            "--",
+            "dune",
+            "build",
+            &format!("{}/callable.exe.o", ocaml_callable_dir),
+        ])
         .status()
         .expect("Dune failed");
     Command::new("rm")

diff --git a/testing/rust-caller/src/lib.rs b/testing/rust-caller/src/lib.rs
@@ -3,6 +3,8 @@
 
 extern crate ocaml_interop;
 
+#[cfg(test)]
+use ocaml_interop::OCamlInt64;
 use ocaml_interop::{ocaml_frame, to_ocaml, OCaml, OCamlBytes, OCamlRuntime, ToOCaml};
 
 mod ocaml {
@@ -255,7 +257,6 @@ fn test_variant_conversion() {
     );
 }
 
-
 #[test]
 #[serial]
 fn test_exception_handling_with_message() {
@@ -269,7 +270,10 @@ fn test_exception_handling_with_message() {
         });
     });
     assert_eq!(
-        result.err().and_then(|err| Some(err.downcast_ref::<String>().unwrap().clone())).unwrap(),
+        result
+            .err()
+            .and_then(|err| Some(err.downcast_ref::<String>().unwrap().clone()))
+            .unwrap(),
         "OCaml exception, message: Some(\"my-error-message\")"
     );
 }
@@ -283,7 +287,10 @@ fn test_exception_handling_without_message() {
         ocaml::raises_nonmessage_exception(cr, &OCaml::unit());
     });
     assert_eq!(
-        result.err().and_then(|err| Some(err.downcast_ref::<String>().unwrap().clone())).unwrap(),
+        result
+            .err()
+            .and_then(|err| Some(err.downcast_ref::<String>().unwrap().clone()))
+            .unwrap(),
         "OCaml exception, message: None"
     );
 }
@@ -297,7 +304,42 @@ fn test_exception_handling_nonblock_exception() {
         ocaml::raises_nonblock_exception(cr, &OCaml::unit());
     });
     assert_eq!(
-        result.err().and_then(|err| Some(err.downcast_ref::<String>().unwrap().clone())).unwrap(),
+        result
+            .err()
+            .and_then(|err| Some(err.downcast_ref::<String>().unwrap().clone()))
+            .unwrap(),
         "OCaml exception, message: None"
     );
-}
+}
+
+#[test]
+#[serial]
+fn test_global_roots() {
+    OCamlRuntime::init_persistent();
+    let mut cr = unsafe { OCamlRuntime::recover_handle() };
+    let crr = &mut cr;
+
+    let i64: OCaml<OCamlInt64> = to_ocaml!(crr, 5);
+    let root = i64.register_global_root();
+    ocaml::gc_compact(crr, &OCaml::unit());
+    root.set(to_ocaml!(crr, 6));
+    ocaml::gc_compact(crr, &OCaml::unit());
+    let i64_bis: i64 = crr.get(root.get_ref()).to_rust();
+    assert_eq!(i64_bis, 6);
+}
+
+#[test]
+#[serial]
+fn test_generational_roots() {
+    OCamlRuntime::init_persistent();
+    let mut cr = unsafe { OCamlRuntime::recover_handle() };
+    let crr = &mut cr;
+
+    let i64: OCaml<OCamlInt64> = to_ocaml!(crr, 5);
+    let root = i64.register_generational_root();
+    ocaml::gc_compact(crr, &OCaml::unit());
+    root.set(to_ocaml!(crr, 6));
+    ocaml::gc_compact(crr, &OCaml::unit());
+    let i64_bis: i64 = crr.get(root.get_ref()).to_rust();
+    assert_eq!(i64_bis, 6);
+}