From d85c8ad5599929bc09c5d387604356261178d4a0 Mon Sep 17 00:00:00 2001
From: Duddino <duddino@duddino.com>
Date: Fri, 5 Apr 2024 20:50:54 +0200
Subject: [PATCH] Add getSaplingRoot and reloadFromCheckpoint

---
 js/pivx_shield.ts | 24 ++++++++++++++++++++++++
 src/utils.rs      | 21 +++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/js/pivx_shield.ts b/js/pivx_shield.ts
index 86b49be..bf569b2 100644
--- a/js/pivx_shield.ts
+++ b/js/pivx_shield.ts
@@ -464,6 +464,30 @@ export class PIVXShield {
   getLastSyncedBlock() {
     return this.lastProcessedBlock;
   }
+
+  /**
+   * @returns sapling root
+   */
+  async getSaplingRoot(): Promise<string> {
+    return await this.callWorker<string>(
+      "get_sapling_root",
+      this.commitmentTree,
+    );
+  }
+
+  /**
+   * Reloads from checkpoint. Needs to be resynced to use
+   */
+  async reloadFromCheckpoint(checkpointBlock: number): Promise<void> {
+    const [effectiveHeight, commitmentTree] = await this.callWorker<
+      [number, string]
+    >("get_closest_checkpoint", checkpointBlock, this.isTestnet);
+    this.commitmentTree = commitmentTree;
+    this.lastProcessedBlock = effectiveHeight;
+    this.unspentNotes = [];
+    this.pendingSpentNotes = new Map();
+    this.pendingUnspentNotes = new Map();
+  }
 }
 
 export interface UTXO {
diff --git a/src/utils.rs b/src/utils.rs
index b1d7929..d7e2ab3 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,3 +1,11 @@
+use std::io::Cursor;
+
+use pivx_primitives::{
+    merkle_tree::{CommitmentTree, HashSer},
+    sapling::Node,
+};
+pub use wasm_bindgen::prelude::*;
+
 pub fn set_panic_hook() {
     // When the `console_error_panic_hook` feature is enabled, we can call the
     // `set_panic_hook` function at least once during initialization, and then
@@ -8,3 +16,16 @@ pub fn set_panic_hook() {
     #[cfg(feature = "console_error_panic_hook")]
     console_error_panic_hook::set_once();
 }
+
+#[wasm_bindgen]
+pub fn get_sapling_root(tree_hex: &str) -> Result<JsValue, JsValue> {
+    let buff = Cursor::new(
+        hex::decode(tree_hex).map_err(|_| "Cannot decode commitment tree from hexadecimal")?,
+    );
+    let tree = CommitmentTree::<Node>::read(buff).map_err(|_| "Cannot decode commitment tree!")?;
+    let mut root = Vec::new();
+    tree.root()
+        .write(&mut root)
+        .map_err(|_| "Cannot write sapling root")?;
+    Ok(JsValue::from_str(&hex::encode(root)))
+}