Add global and generational roots

src/lib.rs

@@ -300,7 +300,7 @@ pub use crate::boxroot::BoxRoot;
 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,
 };

src/memory.rs

@@ -7,12 +7,115 @@ use crate::{
     runtime::OCamlRuntime,
     value::OCaml,
 };
-use core::{cell::UnsafeCell, marker::PhantomData};
+use core::{
+    cell::{Cell, UnsafeCell},
+    marker::PhantomData,
+    pin::Pin,
+};
 pub use ocaml_sys::{caml_alloc, store_field};
 use ocaml_sys::{
     caml_alloc_string, caml_alloc_tuple, caml_copy_double, caml_copy_int32, caml_copy_int64, string_val,
 };
 
+/// 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>,

src/value.rs

@@ -1,6 +1,7 @@
 // Copyright (c) SimpleStaking and Tezedge Contributors
 // SPDX-License-Identifier: MIT
 
+use crate::memory::{OCamlGenerationalRoot, OCamlGlobalRoot};
 use crate::{
     boxroot::BoxRoot, error::OCamlFixnumConversionError, memory::OCamlCell, mlvalues::*, FromOCaml,
     OCamlRef, OCamlRuntime,
@@ -113,6 +114,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, ()> {

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")

testing/rust-caller/src/lib.rs

@@ -3,6 +3,8 @@
 
 extern crate ocaml_interop;
 
+#[cfg(test)]
+use ocaml_interop::OCamlInt64;
 use ocaml_interop::{OCaml, OCamlBytes, OCamlRuntime, ToOCaml};
 
 mod ocaml {
@@ -330,3 +332,35 @@ fn test_exception_handling_nonblock_exception() {
         "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> = 5.to_ocaml(crr);
+    let root = i64.register_global_root();
+    ocaml::gc_compact(crr, &OCaml::unit());
+    root.set(6.to_ocaml(crr));
+    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> = 5.to_ocaml(crr);
+    let root = i64.register_generational_root();
+    ocaml::gc_compact(crr, &OCaml::unit());
+    root.set(6.to_ocaml(crr));
+    ocaml::gc_compact(crr, &OCaml::unit());
+    let i64_bis: i64 = crr.get(root.get_ref()).to_rust();
+    assert_eq!(i64_bis, 6);
+}