drk: properly handle resets

bin/drk/src/cache.rs

@@ -80,19 +80,26 @@ impl Cache {
         })
     }
 
-    /// Fetch given block height numbers from the store's state inverse
-    /// diffs tree. The function will fail if a block height number was
-    /// not found.
-    pub fn get_state_inverse_diff(&self, heights: &[u32]) -> Result<Vec<SledDbOverlayStateDiff>> {
-        let mut ret = Vec::with_capacity(heights.len());
-        for height in heights {
-            match self.state_inverse_diff.get(height.to_be_bytes())? {
-                Some(found) => ret.push(deserialize(&found)?),
-                None => return Err(Error::BlockStateInverseDiffNotFound(*height)),
-            };
-        }
+    /// Insert a `u32` and a block inverse diff into store's inverse
+    /// diffs tree. The block height is used as the key, and the
+    /// serialized database inverse diff is used as value.
+    pub fn insert_state_inverse_diff(
+        &self,
+        height: &u32,
+        diff: &SledDbOverlayStateDiff,
+    ) -> Result<()> {
+        self.state_inverse_diff.insert(height.to_be_bytes(), serialize(diff))?;
+        Ok(())
+    }
 
-        Ok(ret)
+    /// Fetch given block height number from the store's state inverse
+    /// diffs tree. The function will fail if the block height number
+    /// was not found.
+    pub fn get_state_inverse_diff(&self, height: &u32) -> Result<SledDbOverlayStateDiff> {
+        match self.state_inverse_diff.get(height.to_be_bytes())? {
+            Some(found) => Ok(deserialize(&found)?),
+            None => Err(Error::BlockStateInverseDiffNotFound(*height)),
+        }
     }
 }
 
@@ -132,18 +139,6 @@ impl CacheOverlay {
         Ok(())
     }
 
-    /// Insert a `u32` and a block inverse diff into overlay's inverse
-    /// diffs tree. The block height is used as the key, and the
-    /// serialized database inverse diff is used as value.
-    pub fn insert_state_inverse_diff(
-        &mut self,
-        height: &u32,
-        diff: &SledDbOverlayStateDiff,
-    ) -> Result<()> {
-        self.0.insert(SLED_STATE_INVERSE_DIFF_TREE, &height.to_be_bytes(), &serialize(diff))?;
-        Ok(())
-    }
-
     /// Insert a bytes slice and a merkle tree into overlay's merkle
     /// trees tree. The provided bytes slice is used as the key, and
     /// the serialized merkle tree is used as value.

bin/drk/src/rpc.rs

@@ -410,13 +410,6 @@ impl Drk {
         // Grab the overlay current diff
         let diff = scan_cache.money_smt.store.overlay.0.diff(&[])?;
 
-        // Insert the state inverse diff record
-        scan_cache
-            .money_smt
-            .store
-            .overlay
-            .insert_state_inverse_diff(&block.header.height, &diff.inverse())?;
-
         // Apply the overlay current changes
         scan_cache
             .money_smt
@@ -425,6 +418,12 @@ impl Drk {
             .0
             .apply_diff(&scan_cache.money_smt.store.overlay.0.diff(&[])?)?;
 
+        // Insert the state inverse diff record
+        self.cache.insert_state_inverse_diff(&block.header.height, &diff.inverse())?;
+
+        // Flush sled
+        self.cache.sled_db.flush()?;
+
         // Update wallet transactions records
         if let Err(e) =
             self.put_tx_history_records(&wallet_txs, "Confirmed", Some(block.header.height)).await

bin/drk/src/scanned_blocks.rs

@@ -66,22 +66,28 @@ impl Drk {
             return Err(WalletDbError::QueryExecutionFailed);
         };
         let Some((key, value)) = query_result else { return Ok((0, String::from("-"))) };
-        let Ok(key) = deserialize(&key) else {
-            return Err(WalletDbError::ParseColumnValueError);
+        let key: [u8; 4] = match key.as_ref().try_into() {
+            Ok(k) => k,
+            Err(_) => return Err(WalletDbError::ParseColumnValueError),
         };
+        let key = u32::from_be_bytes(key);
         let Ok(value) = deserialize(&value) else {
             return Err(WalletDbError::ParseColumnValueError);
         };
         Ok((key, value))
     }
 
-    /// Reset the scanned blocks information records in the wallet.
+    /// Reset the scanned blocks information records in the cache.
     pub fn reset_scanned_blocks(&self) -> WalletDbResult<()> {
         println!("Resetting scanned blocks");
         if let Err(e) = self.cache.scanned_blocks.clear() {
             println!("[reset_scanned_blocks] Resetting scanned blocks tree failed: {e:?}");
             return Err(WalletDbError::GenericError)
         }
+        if let Err(e) = self.cache.state_inverse_diff.clear() {
+            println!("[reset_scanned_blocks] Resetting state inverse diffs tree failed: {e:?}");
+            return Err(WalletDbError::GenericError)
+        }
         println!("Successfully reset scanned blocks");
 
         Ok(())
@@ -107,19 +113,6 @@ impl Drk {
             return Ok(())
         }
 
-        // Grab all state inverse diffs until requested height,
-        // going backwards.
-        let heights: Vec<u32> = (height + 1..=last).rev().collect();
-        let inverse_diffs = match self.cache.get_state_inverse_diff(&heights) {
-            Ok(d) => d,
-            Err(e) => {
-                println!(
-                    "[reset_to_height] Retrieving state inverse diffs from cache failed: {e:?}"
-                );
-                return Err(WalletDbError::GenericError)
-            }
-        };
-
         // Create an overlay to apply the reverse diffs
         let mut overlay = match CacheOverlay::new(&self.cache) {
             Ok(o) => o,
@@ -129,8 +122,20 @@ impl Drk {
             }
         };
 
-        // Apply the inverse diffs sequence
-        for inverse_diff in inverse_diffs {
+        // Grab all state inverse diffs until requested height,
+        // going backwards.
+        for height in (height + 1..=last).rev() {
+            let inverse_diff = match self.cache.get_state_inverse_diff(&height) {
+                Ok(d) => d,
+                Err(e) => {
+                    println!(
+                        "[reset_to_height] Retrieving state inverse diff from cache failed: {e:?}"
+                    );
+                    return Err(WalletDbError::GenericError)
+                }
+            };
+
+            // Apply it
             if let Err(e) = overlay.0.add_diff(&inverse_diff) {
                 println!("[reset_to_height] Adding state inverse diff to the cache overlay failed: {e:?}");
                 return Err(WalletDbError::GenericError)
@@ -139,6 +144,16 @@ impl Drk {
                 println!("[reset_to_height] Applying state inverse diff to the cache overlay failed: {e:?}");
                 return Err(WalletDbError::GenericError)
             }
+
+            // Remove it
+            if let Err(e) = self.cache.state_inverse_diff.remove(height.to_be_bytes()) {
+                println!(
+                    "[reset_to_height] Removing state inverse diff from the cache failed: {e:?}"
+                );
+                return Err(WalletDbError::GenericError)
+            }
+
+            // Flush sled
             if let Err(e) = self.cache.sled_db.flush() {
                 println!("[reset_to_height] Flushing cache sled database failed: {e:?}");
                 return Err(WalletDbError::GenericError)