Merge PR use pin unpin collect

Retrieve pin unpin and collect in Env. Add a rewrite
adrien-zinger · Dec 9, 2021 · ff4f5c8 · ff4f5c8
2 parents 8ab6439 + 0a95e06
commit ff4f5c8
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 21 deletions.
diff --git a/src/env.rs b/src/env.rs
@@ -3,15 +3,29 @@ use wasmer::{Function, HostEnvInitError, Instance, LazyInit, Memory, WasmerEnv};
 #[derive(Clone, Default)]
 pub struct Env {
     pub memory: LazyInit<Memory>,
-    pub new: Option<Function>,
+    pub fn_new: Option<Function>,
+    pub fn_pin: Option<Function>,
+    pub fn_unpin: Option<Function>,
+    pub fn_collect: Option<Function>,
 }
 
 impl Env {
-    pub fn new(memory: Memory, new: Option<Function>) -> Env {
-        let mut env = Env::default();
-        env.memory.initialize(memory);
-        env.new = new;
-        env
+    pub fn new(
+        arg_memory: Memory,
+        fn_new: Option<Function>,
+        fn_pin: Option<Function>,
+        fn_unpin: Option<Function>,
+        fn_collect: Option<Function>,
+    ) -> Env {
+        let mut memory = LazyInit::<Memory>::default();
+        memory.initialize(arg_memory);
+        Env {
+            memory,
+            fn_new,
+            fn_pin,
+            fn_unpin,
+            fn_collect,
+        }
     }
 }
 
@@ -23,7 +37,16 @@ impl WasmerEnv for Env {
             .map_err(HostEnvInitError::from)?
             .clone();
         if let Ok(func) = instance.exports.get_function("__new") {
-            self.new = Some(func.clone())
+            self.fn_new = Some(func.clone())
+        }
+        if let Ok(func) = instance.exports.get_function("__pin") {
+            self.fn_pin = Some(func.clone())
+        }
+        if let Ok(func) = instance.exports.get_function("__unpin") {
+            self.fn_unpin = Some(func.clone())
+        }
+        if let Ok(func) = instance.exports.get_function("__collect") {
+            self.fn_collect = Some(func.clone())
         }
         self.memory.initialize(mem);
         Ok(())

diff --git a/src/lib.rs b/src/lib.rs
@@ -16,8 +16,8 @@ pub trait Read<T> {
 
 pub trait Write<T> {
     fn alloc(value: &str, memory: &Env) -> anyhow::Result<Box<Self>>;
-    fn write(&self, value: &str, env: &Env) -> anyhow::Result<()>;
-    fn free(memory: &Env) -> anyhow::Result<()>;
+    fn write(&mut self, value: &str, env: &Env) -> anyhow::Result<Box<Self>>;
+    fn free(self, memory: &Env) -> anyhow::Result<()>;
 }
 
 #[derive(Debug)]

diff --git a/src/string_ptr.rs b/src/string_ptr.rs
@@ -1,9 +1,17 @@
 use super::{Env, Memory, Read, Write};
-use std::convert::TryFrom;
+use std::convert::{TryFrom, TryInto};
 use wasmer::{Array, Value, WasmPtr};
 
 pub type StringPtr = WasmPtr<u16, Array>;
 
+macro_rules! export_asr {
+    ($func_name:ident, $env:expr) => {
+        $env.$func_name
+            .as_ref()
+            .expect("Assembly Script Runtime not exported")
+    };
+}
+
 impl Read<String> for StringPtr {
     fn read(self, memory: &Memory) -> anyhow::Result<String> {
         let size = self.size(memory)?;
@@ -24,10 +32,7 @@ impl Read<String> for StringPtr {
 
 impl Write<String> for StringPtr {
     fn alloc(value: &str, env: &Env) -> anyhow::Result<Box<StringPtr>> {
-        let new = env
-            .new
-            .as_ref()
-            .expect("Assembly Script Runtime ot exported");
+        let new = export_asr!(fn_new, env);
         let size = i32::try_from(value.len())?;
 
         let offset = u32::try_from(
@@ -39,25 +44,51 @@ impl Write<String> for StringPtr {
                 .unwrap(),
         )?;
         write_str(offset, value, env)?;
+
+        // pin
+        let pin = export_asr!(fn_pin, env);
+        pin.call(&[Value::I32(offset.try_into().unwrap())])
+            .expect("Failed to call __pin");
+
         Ok(Box::new(StringPtr::new(offset)))
     }
 
-    fn write(&self, value: &str, env: &Env) -> anyhow::Result<()> {
+    fn write(&mut self, value: &str, env: &Env) -> anyhow::Result<Box<StringPtr>> {
         let prev_size = size(
             self.offset(),
             env.memory.get_ref().expect("Failed to load memory"),
         )?;
         let new_size = u32::try_from(value.len())? << 1;
         if prev_size == new_size {
-            write_str(self.offset(), value, env)?
+            write_str(self.offset(), value, env)?;
+            Ok(Box::new(*self))
         } else {
-            todo!("Remove this and reallocate of bigger or smaller space")
+            // unpin old ptr
+            let unpin = export_asr!(fn_pin, env);
+            unpin
+                .call(&[Value::I32(self.offset().try_into().unwrap())])
+                .expect("Failed to call __unpin");
+
+            // collect
+            let collect = export_asr!(fn_collect, env);
+            collect.call(&[]).expect("failed to call __collect");
+
+            // alloc with new size
+            StringPtr::alloc(value, env)
         }
-        Ok(())
     }
 
-    fn free(_env: &Env) -> anyhow::Result<()> {
-        todo!("Release the memory from this string")
+    fn free(self, env: &Env) -> anyhow::Result<()> {
+        // unpin
+        let unpin = export_asr!(fn_pin, env);
+        unpin
+            .call(&[Value::I32(self.offset().try_into().unwrap())])
+            .expect("Failed to call __unpin");
+
+        // collect
+        let collect = export_asr!(fn_collect, env);
+        collect.call(&[]).expect("failed to call __collect");
+        Ok(())
     }
 }
 

diff --git a/tests/strings.rs b/tests/strings.rs
@@ -52,6 +52,18 @@ fn read_alloc_strings() -> Result<(), Box<dyn Error>> {
             Ok(func) => Some(func.clone()),
             _ => None,
         },
+        match instance.exports.get_function("__pin") {
+            Ok(func) => Some(func.clone()),
+            _ => None,
+        },
+        match instance.exports.get_function("__unpin") {
+            Ok(func) => Some(func.clone()),
+            _ => None,
+        },
+        match instance.exports.get_function("__collect") {
+            Ok(func) => Some(func.clone()),
+            _ => None,
+        },
     );
 
     let get_string = instance
@@ -94,13 +106,25 @@ fn read_write_strings() -> Result<(), Box<dyn Error>> {
             Ok(func) => Some(func.clone()),
             _ => None,
         },
+        match instance.exports.get_function("__pin") {
+            Ok(func) => Some(func.clone()),
+            _ => None,
+        },
+        match instance.exports.get_function("__unpin") {
+            Ok(func) => Some(func.clone()),
+            _ => None,
+        },
+        match instance.exports.get_function("__collect") {
+            Ok(func) => Some(func.clone()),
+            _ => None,
+        },
     );
 
     let get_string = instance
         .exports
         .get_native_function::<(), StringPtr>("getString")?;
 
-    let str_ptr = get_string.call()?;
+    let mut str_ptr = get_string.call()?;
     let string = str_ptr.read(memory)?;
 
     assert_eq!(string, "hello test");