From 718f5dc90debfa44cb698da3b05ef660e8b32534 Mon Sep 17 00:00:00 2001
From: Tony <legendmastertony@gmail.com>
Date: Thu, 3 Oct 2024 16:46:06 +0800
Subject: [PATCH] Allow js to use pre-stored (de)serialize functions

---
 plugins/store/guest-js/index.ts |  8 +++++
 plugins/store/src/error.rs      |  6 ++++
 plugins/store/src/lib.rs        | 55 +++++++++++++++++++++++++++++----
 plugins/store/src/store.rs      |  4 +--
 4 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/plugins/store/guest-js/index.ts b/plugins/store/guest-js/index.ts
index aecf9e2f..cc85c78a 100644
--- a/plugins/store/guest-js/index.ts
+++ b/plugins/store/guest-js/index.ts
@@ -21,6 +21,14 @@ export type StoreOptions = {
    * Auto save on modification with debounce duration in milliseconds, it's 100ms by default, pass in `false` to disable it
    */
   autoSave?: boolean | number
+  /**
+   * Name of a serialize function registered in the rust side plugin builder
+   */
+  serializeFnName?: string
+  /**
+   * Name of a deserialize function registered in the rust side plugin builder
+   */
+  deserializeFnName?: string
 }
 
 /**
diff --git a/plugins/store/src/error.rs b/plugins/store/src/error.rs
index 612edcfd..77526fbc 100644
--- a/plugins/store/src/error.rs
+++ b/plugins/store/src/error.rs
@@ -24,6 +24,12 @@ pub enum Error {
     /// Store already exists
     #[error("Store at \"{0}\" already exists")]
     AlreadyExists(PathBuf),
+    /// Serialize function not found
+    #[error("Serialize Function \"{0}\" not found")]
+    SerializeFunctionNotFound(String),
+    /// Deserialize function not found
+    #[error("Deserialize Function \"{0}\" not found")]
+    DeserializeFunctionNotFound(String),
     /// Some Tauri API failed
     #[error(transparent)]
     Tauri(#[from] tauri::Error),
diff --git a/plugins/store/src/lib.rs b/plugins/store/src/lib.rs
index 40f72aff..6e01061b 100644
--- a/plugins/store/src/lib.rs
+++ b/plugins/store/src/lib.rs
@@ -21,10 +21,10 @@ use std::{
     sync::{Arc, Mutex},
     time::Duration,
 };
-pub use store::{Store, StoreBuilder, StoreInner};
+pub use store::{DeserializeFn, SerializeFn, Store, StoreBuilder, StoreInner};
 use tauri::{
     plugin::{self, TauriPlugin},
-    AppHandle, Manager, ResourceId, RunEvent, Runtime,
+    AppHandle, Manager, ResourceId, RunEvent, Runtime, State,
 };
 
 mod error;
@@ -41,6 +41,8 @@ struct ChangePayload<'a> {
 
 pub struct StoreCollection {
     stores: Mutex<HashMap<PathBuf, ResourceId>>,
+    serialize_fns: HashMap<String, SerializeFn>,
+    deserialize_fns: HashMap<String, DeserializeFn>,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -53,10 +55,14 @@ enum AutoSave {
 #[tauri::command]
 async fn create_store<R: Runtime>(
     app: AppHandle<R>,
+    store_collection: State<'_, StoreCollection>,
     path: PathBuf,
     auto_save: Option<AutoSave>,
+    serialize_fn_name: Option<String>,
+    deserialize_fn_name: Option<String>,
 ) -> Result<ResourceId> {
     let mut builder = app.store_builder(path.clone());
+
     if let Some(auto_save) = auto_save {
         match auto_save {
             AutoSave::DebounceDuration(duration) => {
@@ -68,15 +74,34 @@ async fn create_store<R: Runtime>(
             _ => {}
         }
     }
+
+    if let Some(serialize_fn_name) = serialize_fn_name {
+        let serialize_fn = store_collection
+            .serialize_fns
+            .get(&serialize_fn_name)
+            .ok_or_else(|| crate::Error::SerializeFunctionNotFound(serialize_fn_name))?;
+        builder = builder.serialize(*serialize_fn);
+    }
+
+    if let Some(deserialize_fn_name) = deserialize_fn_name {
+        let deserialize_fn = store_collection
+            .deserialize_fns
+            .get(&deserialize_fn_name)
+            .ok_or_else(|| crate::Error::DeserializeFunctionNotFound(deserialize_fn_name))?;
+        builder = builder.deserialize(*deserialize_fn);
+    }
+
     let (_, rid) = builder.build_inner()?;
     Ok(rid)
 }
 
 #[tauri::command]
-async fn get_store<R: Runtime>(app: AppHandle<R>, path: PathBuf) -> Option<ResourceId> {
-    let collection = app.state::<StoreCollection>();
-    let stores = collection.stores.lock().unwrap();
-    stores.get(&path).copied()
+async fn get_store(
+    store_collection: State<'_, StoreCollection>,
+    path: PathBuf,
+) -> Result<Option<ResourceId>> {
+    let stores = store_collection.stores.lock().unwrap();
+    Ok(stores.get(&path).copied())
 }
 
 #[tauri::command]
@@ -202,12 +227,16 @@ impl<R: Runtime, T: Manager<R>> StoreExt<R> for T {
 
 pub struct Builder<R: Runtime> {
     phantom_data: PhantomData<R>,
+    serialize_fns: HashMap<String, SerializeFn>,
+    deserialize_fns: HashMap<String, DeserializeFn>,
 }
 
 impl<R: Runtime> Default for Builder<R> {
     fn default() -> Self {
         Self {
             phantom_data: Default::default(),
+            serialize_fns: Default::default(),
+            deserialize_fns: Default::default(),
         }
     }
 }
@@ -217,6 +246,18 @@ impl<R: Runtime> Builder<R> {
         Self::default()
     }
 
+    /// Register a serialize function to access it from the JavaScript side
+    pub fn register_serialize_fn(mut self, name: String, serialize_fn: SerializeFn) -> Self {
+        self.serialize_fns.insert(name, serialize_fn);
+        self
+    }
+
+    /// Register a deserialize function to access it from the JavaScript side
+    pub fn register_deserialize_fn(mut self, name: String, deserialize_fn: DeserializeFn) -> Self {
+        self.deserialize_fns.insert(name, deserialize_fn);
+        self
+    }
+
     /// Builds the plugin.
     ///
     /// # Examples
@@ -250,6 +291,8 @@ impl<R: Runtime> Builder<R> {
             .setup(move |app_handle, _api| {
                 app_handle.manage(StoreCollection {
                     stores: Mutex::new(HashMap::new()),
+                    serialize_fns: self.serialize_fns,
+                    deserialize_fns: self.deserialize_fns,
                 });
                 Ok(())
             })
diff --git a/plugins/store/src/store.rs b/plugins/store/src/store.rs
index b2c4c23d..aa5117bd 100644
--- a/plugins/store/src/store.rs
+++ b/plugins/store/src/store.rs
@@ -19,9 +19,9 @@ use tokio::{
     time::sleep,
 };
 
-type SerializeFn =
+pub type SerializeFn =
     fn(&HashMap<String, JsonValue>) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>>;
-pub(crate) type DeserializeFn =
+pub type DeserializeFn =
     fn(&[u8]) -> Result<HashMap<String, JsonValue>, Box<dyn std::error::Error + Send + Sync>>;
 
 fn default_serialize(