State rollback and wipe functions (including wasm lib)

libzkbob-rs-wasm/src/client/mod.rs

@@ -523,4 +523,14 @@ impl UserAccount {
 
         serde_wasm_bindgen::to_value(&data).unwrap()
     }
+
+    #[wasm_bindgen(js_name = "rollbackState")]
+    pub fn rollback_state(&self, rollback_index: u64) -> u64 {
+        self.inner.borrow_mut().state.rollback(rollback_index)
+    }
+
+    #[wasm_bindgen(js_name = "wipeState")]
+    pub fn wipe_state(&self) {
+        self.inner.borrow_mut().state.wipe();
+    }
 }

libzkbob-rs/src/client/state.rs

@@ -257,4 +257,44 @@ where
 
         note_balance
     }
+
+    // rollback current state, return updated next_index
+    pub fn rollback(&mut self, rollback_index: u64) -> u64 {
+        if rollback_index > self.tree.next_index() {
+            return self.tree.next_index();
+        }
+
+        let rollback_index = rollback_index & ((1 << constants::OUTPLUSONELOG) - 1);
+        self.txs.remove_from(rollback_index);
+        for (index, tx) in self.txs.iter() {
+            match tx {
+                Transaction::Account(acc) => {
+                    if index >= self.latest_account_index.unwrap_or(0) {
+                        self.latest_account_index = Some(index);
+                        self.latest_account = Some(acc);
+                    }
+                }
+                Transaction::Note(_) => {
+                    if index >= self.latest_note_index {
+                        self.latest_note_index = index;
+                    }
+                }
+            }
+        }
+
+        if self.tree.rollback(rollback_index).is_none() {
+            self.wipe();
+        }
+
+        self.tree.next_index()
+    }
+
+    pub fn wipe(&mut self) {
+        self.txs.remove_all();
+        self.latest_account_index = None;
+        self.latest_account = None;
+        self.latest_note_index = 0;
+
+    }
+
 }

libzkbob-rs/src/sparse_array.rs

@@ -112,6 +112,23 @@ where
         self.db.write(batch).unwrap();
     }
 
+    pub fn remove_from(&self, from_index: u64) {
+        let mut batch = self.db.transaction();
+        for (index, _) in self.iter() {
+            if index >= from_index {
+                let key = index.to_be_bytes();
+                batch.delete(0, &key);
+            }
+        }
+        self.db.write(batch).unwrap();
+    }
+
+    pub fn remove_all(&self) {
+        let mut batch = self.db.transaction();
+        batch.delete_prefix(0, &[]);
+        self.db.write(batch).unwrap();
+    }
+
     // FIXME: Crazy inefficient, replace or improve kvdb
     pub fn count(&self) -> usize {
         self.db.iter(0).count()
@@ -176,4 +193,21 @@ mod tests {
         assert_eq!(a.iter_slice(2..=412345).count(), 2, "from 2");
         assert_eq!(a.iter_slice(2..=412344).count(), 1, "from 2 except last");
     }
+
+    #[test]
+    fn test_sparse_array_remove() {
+        let a = SparseArray::new_test();
+        a.set(1, &1u32);
+        a.set(3, &2);
+        a.set(10, &3);
+        a.set(20, &4);
+        a.set(25, &5);
+        a.set(100, &6);
+
+        a.remove_from(10);
+        assert_eq!(a.iter_slice(0..=100).count(), 3);
+
+        a.remove_all();
+        assert_eq!(a.iter_slice(0..=100).count(), 0);
+    }
 }