[3/n][vm-rewrite][Move] Update native context extensions to support Cloning

external-crates/move/crates/move-vm-runtime/src/natives/extensions.rs

@@ -3,7 +3,34 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use better_any::{Tid, TidAble, TidExt};
-use std::{any::TypeId, collections::HashMap};
+use std::{any::TypeId, cell::RefCell, collections::HashMap, ops::Deref, rc::Rc};
+
+/// A helper wrapper around a `Tid`able type that encapsulates interior mutability in a single-threaded
+/// manner.
+///
+/// Note that this is _not_ threadsafe. If you need threadsafe access to the `T` you will need to
+/// handle that within `T'`s type (just like in the previous implementation of the
+/// `NativeContextExtensions`).
+#[derive(Tid)]
+pub struct NativeContextMut<'a, T: Tid<'a>>(pub RefCell<T>, std::marker::PhantomData<&'a ()>);
+
+impl<'a, T: Tid<'a>> NativeContextMut<'a, T> {
+    /// Create a new `NativeContextMut` value with the given value.
+    pub fn new(t: T) -> Self {
+        NativeContextMut(RefCell::new(t), std::marker::PhantomData)
+    }
+
+    pub fn into_inner(self) -> T {
+        self.0.into_inner()
+    }
+}
+
+impl<'a, T: Tid<'a>> Deref for NativeContextMut<'a, T> {
+    type Target = RefCell<T>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
 
 /// A data type to represent a heterogeneous collection of extensions which are available to
 /// native functions. A value to this is passed into the session function execution.
@@ -13,15 +40,15 @@ use std::{any::TypeId, collections::HashMap};
 /// avoids that extensions need to have `'static` lifetime, which `Any` requires. In order to make a
 /// struct suitable to be a 'Tid', use `#[derive(Tid)]` in the struct declaration. (See also
 /// tests at the end of this module.)
-#[derive(Default)]
+#[derive(Default, Clone)]
 pub struct NativeContextExtensions<'a> {
-    map: HashMap<TypeId, Box<dyn Tid<'a>>>,
+    map: HashMap<TypeId, Rc<dyn Tid<'a>>>,
 }
 
 impl<'a> NativeContextExtensions<'a> {
     pub fn add<T: TidAble<'a>>(&mut self, ext: T) {
         assert!(
-            self.map.insert(T::id(), Box::new(ext)).is_none(),
+            self.map.insert(T::id(), Rc::new(ext)).is_none(),
             "multiple extensions of the same type not allowed"
         )
     }
@@ -35,32 +62,23 @@ impl<'a> NativeContextExtensions<'a> {
             .unwrap()
     }
 
-    pub fn get_mut<T: TidAble<'a>>(&mut self) -> &mut T {
-        self.map
-            .get_mut(&T::id())
-            .expect("extension unknown")
-            .as_mut()
-            .downcast_mut::<T>()
-            .unwrap()
-    }
-
-    pub fn remove<T: TidAble<'a>>(&mut self) -> T {
+    pub fn remove<T: TidAble<'a>>(&mut self) -> Rc<T> {
         // can't use expect below because it requires `T: Debug`.
         match self
             .map
             .remove(&T::id())
             .expect("extension unknown")
-            .downcast_box::<T>()
+            .downcast_rc::<T>()
         {
-            Ok(val) => *val,
+            Ok(val) => val,
             Err(_) => panic!("downcast error"),
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::natives::extensions::NativeContextExtensions;
+    use super::*;
     use better_any::{Tid, TidAble};
 
     #[derive(Tid)]
@@ -73,11 +91,11 @@ mod tests {
         let mut v: u64 = 23;
         let e = Ext { a: &mut v };
         let mut exts = NativeContextExtensions::default();
-        exts.add(e);
-        *exts.get_mut::<Ext>().a += 1;
-        assert_eq!(*exts.get_mut::<Ext>().a, 24);
-        *exts.get_mut::<Ext>().a += 1;
-        let e1 = exts.remove::<Ext>();
-        assert_eq!(*e1.a, 25)
+        exts.add(NativeContextMut::new(e));
+        *exts.get::<NativeContextMut<Ext>>().borrow_mut().a += 1;
+        assert_eq!(*exts.get::<NativeContextMut<Ext>>().borrow_mut().a, 24);
+        *exts.get::<NativeContextMut<Ext>>().borrow_mut().a += 1;
+        let e1 = exts.get::<NativeContextMut<Ext>>();
+        assert_eq!(*e1.borrow().a, 25);
     }
 }