fix(drive): rebalance shielded credit pool subtree keys by access frequency

packages/rs-drive-abci/src/execution/platform_events/protocol_upgrade/perform_events_on_first_block_of_protocol_change/v0/mod.rs

@@ -625,8 +625,15 @@ impl<C> Platform<C> {
             &platform_version.drive,
         )?;
 
-        // Notes tree (CommitmentTree = CountTree items + Sinsemilla Frontier):
-        // [AddressBalances, "s"] / [1]
+        // The four child inserts below are ordered breadth-first to match the
+        // intended balanced shape of the parent Merk tree (see the layout
+        // diagram in `drive::drive::shielded::paths`): root first, then both
+        // depth-1 children, then the depth-2 leaf. AVL rebalancing is
+        // order-sensitive, so this ordering is what actually places
+        // `SHIELDED_NOTES_KEY` at the root and the spend-path keys at depth 1.
+
+        // Level 0 (root): notes tree (CommitmentTree = CountTree items + Sinsemilla Frontier)
+        // [AddressBalances, "s"] / [128]
         let shielded_pool_path = shielded_credit_pool_path();
         self.drive.grove_insert_if_not_exists(
             (&shielded_pool_path).into(),
@@ -637,7 +644,8 @@ impl<C> Platform<C> {
             &platform_version.drive,
         )?;
 
-        // Nullifiers tree (ProvableCountTree): [AddressBalances, "s"] / [2]
+        // Level 1 (left): nullifiers tree (ProvableCountTree)
+        // [AddressBalances, "s"] / [64]
         self.drive.grove_insert_if_not_exists(
             (&shielded_pool_path).into(),
             &[SHIELDED_NULLIFIERS_KEY],
@@ -647,22 +655,23 @@ impl<C> Platform<C> {
             &platform_version.drive,
         )?;
 
-        // Total balance SumItem(0): [AddressBalances, "s"] / [5]
+        // Level 1 (right): anchors tree (NormalTree) — anchor_bytes → block_height_be
+        // [AddressBalances, "s"] / [192]
         self.drive.grove_insert_if_not_exists(
             (&shielded_pool_path).into(),
-            &[SHIELDED_TOTAL_BALANCE_KEY],
-            Element::new_sum_item(0),
+            &[SHIELDED_ANCHORS_IN_POOL_KEY],
+            Element::empty_tree(),
             Some(transaction),
             None,
             &platform_version.drive,
         )?;
 
-        // Anchors tree (NormalTree) inside pool: [AddressBalances, "s"] / [6]
-        // Stores block_height_be → anchor_bytes
+        // Level 2: total balance SumItem(0)
+        // [AddressBalances, "s"] / [32]
         self.drive.grove_insert_if_not_exists(
             (&shielded_pool_path).into(),
-            &[SHIELDED_ANCHORS_IN_POOL_KEY],
-            Element::empty_tree(),
+            &[SHIELDED_TOTAL_BALANCE_KEY],
+            Element::new_sum_item(0),
             Some(transaction),
             None,
             &platform_version.drive,

packages/rs-drive-abci/src/execution/validation/state_transition/state_transitions/test_helpers.rs

@@ -462,7 +462,7 @@ pub fn process_transition(
 }
 
 /// Insert a fake anchor into the shielded anchors tree via GroveDB.
-/// Anchors are stored as anchor_bytes → block_height_be in [AddressBalances, "s", [6]].
+/// Anchors are stored as anchor_bytes → block_height_be in [AddressBalances, "s", [192]].
 pub fn insert_anchor_into_state(platform: &TempPlatform<MockCoreRPCLike>, anchor: &[u8; 32]) {
     let platform_version = PlatformVersion::latest();
     let grove_version = &platform_version.drive.grove_version;

packages/rs-drive-abci/src/query/shielded/most_recent_anchor/v0/mod.rs

@@ -18,7 +18,7 @@ use drive::util::grove_operations::GroveDBToUse;
 impl<C> Platform<C> {
     /// Answer `getMostRecentShieldedAnchor` by reading the latest
     /// entry from the anchors-by-height index
-    /// (`[..., "s", [8]]`) — a `limit 1` reverse query. The anchor
+    /// (`[..., "s", [96]]`) — a `limit 1` reverse query. The anchor
     /// is the value at the highest block-height key.
     ///
     /// Returns `[0; 32]` (mapped to `None` by the response decoder

packages/rs-drive/src/drive/initialization/v3/mod.rs

@@ -68,47 +68,57 @@ impl Drive {
     }
 
     /// Adds shielded pool batch operations for initialization.
+    ///
+    /// The eight subtree inserts are ordered breadth-first to match the
+    /// intended balanced shape of the parent Merk tree (see the layout
+    /// diagram in `crate::drive::shielded::paths`): root first, then both
+    /// depth-1 children, then the four depth-2 children, then the depth-3
+    /// leaf. AVL rebalancing is order-sensitive, so this ordering is what
+    /// actually places `SHIELDED_NOTES_KEY` at the root and the spend-path
+    /// keys at depth 1.
     pub(in crate::drive::initialization) fn initial_state_structure_shielded_pool_operations(
         &self,
         batch: &mut GroveDbOpBatch,
     ) -> Result<(), Error> {
-        // 1. Shielded credit pool SumTree under AddressBalances
+        // Parent: shielded credit pool SumTree under AddressBalances. Must be
+        // inserted before any of its children so the subtree exists.
         batch.add_insert(
             vec![vec![RootTree::AddressBalances as u8]],
             vec![SHIELDED_CREDIT_POOL_KEY_U8],
             Element::empty_sum_tree(),
         );
 
-        // 2. Notes tree (CommitmentTree = CountTree items + Sinsemilla Frontier)
+        // Level 0 (root): notes tree (CommitmentTree = CountTree items + Sinsemilla Frontier)
         batch.add_insert(
             shielded_credit_pool_path_vec(),
             vec![SHIELDED_NOTES_KEY],
             Element::empty_commitment_tree(SHIELDED_NOTES_CHUNK_POWER)?,
         );
 
-        // 3. Nullifiers tree (ProvableCountTree)
+        // Level 1 (left): nullifiers tree (ProvableCountTree) — checked on every spend.
         batch.add_insert(
             shielded_credit_pool_path_vec(),
             vec![SHIELDED_NULLIFIERS_KEY],
             Element::empty_provable_count_tree(),
         );
 
-        // 4. Total balance SumItem(0)
+        // Level 1 (right): anchors tree (NormalTree) — checked on every spend.
+        // Stores anchor_bytes → block_height_be.
         batch.add_insert(
             shielded_credit_pool_path_vec(),
-            vec![SHIELDED_TOTAL_BALANCE_KEY],
-            Element::new_sum_item(0),
+            vec![SHIELDED_ANCHORS_IN_POOL_KEY],
+            Element::empty_tree(),
         );
 
-        // 5. Anchors tree (NormalTree) inside pool: anchor_bytes → block_height_be
+        // Level 2: total balance SumItem(0).
         batch.add_insert(
             shielded_credit_pool_path_vec(),
-            vec![SHIELDED_ANCHORS_IN_POOL_KEY],
-            Element::empty_tree(),
+            vec![SHIELDED_TOTAL_BALANCE_KEY],
+            Element::new_sum_item(0),
         );
 
-        // 5b. Anchors-by-height tree (NormalTree): block_height_be → anchor_bytes.
-        // Reverse index for pruning old anchors by height range. Also the
+        // Level 2: anchors-by-height tree (NormalTree) — block_height_be → anchor_bytes.
+        // Reverse index for pruning old anchors by height range and the
         // canonical source of the most-recent anchor (read via `limit 1`
         // reverse query) — there is no separate "most recent" slot; key 7
         // was retired because the duplicate state could desync from the
@@ -119,23 +129,23 @@ impl Drive {
             Element::empty_tree(),
         );
 
-        // 6. Per-block nullifiers CountSumTree under shielded credit pool.
+        // Level 2: per-block recent-nullifiers CountSumTree.
         // Each item is an ItemWithSumItem (serialized Vec<[u8;32]> + nullifier count as sum).
         batch.add_insert(
             shielded_credit_pool_path_vec(),
             vec![SHIELDED_RECENT_NULLIFIERS_KEY_U8],
             Element::empty_count_sum_tree(),
         );
 
-        // 7. Compacted nullifiers NormalTree under shielded credit pool.
-        // Key: (start_block, end_block) as 16 bytes, Value: serialized Vec<[u8;32]>
+        // Level 2: compacted nullifiers NormalTree.
+        // Key: (start_block, end_block) as 16 bytes, Value: serialized Vec<[u8;32]>.
         batch.add_insert(
             shielded_credit_pool_path_vec(),
             vec![SHIELDED_COMPACTED_NULLIFIERS_KEY_U8],
             Element::empty_tree(),
         );
 
-        // 8. Nullifiers expiration time NormalTree under shielded credit pool.
+        // Level 3: nullifiers-expiration-time NormalTree (deepest leaf).
         batch.add_insert(
             shielded_credit_pool_path_vec(),
             vec![SHIELDED_NULLIFIERS_EXPIRATION_TIME_KEY_U8],

packages/rs-drive/src/drive/shielded/estimated_costs.rs

@@ -105,7 +105,7 @@ impl Drive {
             },
         );
 
-        // Notes tree: [AddressBalances, "s", 1]
+        // Notes tree: [AddressBalances, "s", 128]
         // CommitmentTree - stores notes (cmx||encrypted_note items + Sinsemilla frontier)
         estimated_costs_only_with_layer_info.insert(
             KeyInfoPath::from_known_path(shielded_credit_pool_notes_path()),
@@ -116,7 +116,7 @@ impl Drive {
             },
         );
 
-        // Nullifiers tree: [AddressBalances, "s", 2]
+        // Nullifiers tree: [AddressBalances, "s", 64]
         // ProvableCountTree - stores spent nullifiers (32-byte key -> empty item)
         estimated_costs_only_with_layer_info.insert(
             KeyInfoPath::from_known_path(shielded_credit_pool_nullifiers_path()),
@@ -127,7 +127,7 @@ impl Drive {
             },
         );
 
-        // Anchors tree: [AddressBalances, "s", 6]
+        // Anchors tree: [AddressBalances, "s", 192]
         // NormalTree - stores anchor_bytes -> block_height_be
         estimated_costs_only_with_layer_info.insert(
             KeyInfoPath::from_known_path(shielded_credit_pool_anchors_path()),
@@ -138,7 +138,7 @@ impl Drive {
             },
         );
 
-        // Anchors-by-height tree: [AddressBalances, "s", 8]
+        // Anchors-by-height tree: [AddressBalances, "s", 96]
         // NormalTree - stores block_height_be -> anchor_bytes (reverse index for pruning)
         estimated_costs_only_with_layer_info.insert(
             KeyInfoPath::from_known_path(shielded_credit_pool_anchors_by_height_path()),

packages/rs-drive/src/drive/shielded/has_anchor/mod.rs

@@ -11,7 +11,7 @@ impl Drive {
     /// Checks whether a shielded pool anchor exists in the anchors tree.
     ///
     /// Anchors are stored as `anchor_bytes -> block_height_be` in
-    /// `[AddressBalances, "s", [6]]`. Uses O(1) key lookup.
+    /// `[AddressBalances, "s", [192]]`. Uses O(1) key lookup.
     ///
     /// # Parameters
     /// - `anchor`: The 32-byte anchor to look up

packages/rs-drive/src/drive/shielded/has_anchor/v0/mod.rs

@@ -10,7 +10,7 @@ impl Drive {
     /// Version 0 implementation of checking whether a shielded anchor exists.
     ///
     /// Performs an O(1) key lookup in the anchors tree at
-    /// `[AddressBalances, "s", [6]]`.
+    /// `[AddressBalances, "s", [192]]`.
     pub(in crate::drive) fn has_shielded_anchor_v0(
         &self,
         anchor: &[u8; 32],

packages/rs-drive/src/drive/shielded/has_nullifier/mod.rs

@@ -11,7 +11,7 @@ impl Drive {
     /// Checks whether a nullifier has already been spent in the shielded pool.
     ///
     /// Nullifiers are stored in the nullifiers tree at
-    /// `[AddressBalances, "s", [2]]`. Uses O(1) key lookup.
+    /// `[AddressBalances, "s", [64]]`. Uses O(1) key lookup.
     ///
     /// # Parameters
     /// - `nullifier`: The 32-byte nullifier to look up

packages/rs-drive/src/drive/shielded/has_nullifier/v0/mod.rs

@@ -10,7 +10,7 @@ impl Drive {
     /// Version 0 implementation of checking whether a nullifier exists.
     ///
     /// Performs an O(1) key lookup in the nullifiers tree at
-    /// `[AddressBalances, "s", [2]]`.
+    /// `[AddressBalances, "s", [64]]`.
     pub(in crate::drive) fn has_nullifier_v0(
         &self,
         nullifier: &[u8; 32],

packages/rs-drive/src/drive/shielded/notes_count/v0/mod.rs

@@ -9,7 +9,7 @@ impl Drive {
     /// Version 0 implementation of counting shielded pool notes.
     ///
     /// Returns the total number of items in the CommitmentTree at
-    /// `[AddressBalances, "s", [1]]`.
+    /// `[AddressBalances, "s", [128]]`.
     pub(in crate::drive) fn shielded_pool_notes_count_v0(
         &self,
         transaction: TransactionArg,

packages/rs-drive/src/drive/shielded/nullifiers/queries.rs

@@ -1,25 +1,34 @@
 use crate::drive::shielded::paths::SHIELDED_CREDIT_POOL_KEY;
 use crate::drive::RootTree;
 
-/// The subtree key for per-block nullifiers storage (CountSumTree)
-pub const SHIELDED_RECENT_NULLIFIERS_KEY: &[u8; 1] = b"n";
+// Byte positions chosen to balance the parent shielded-pool Merk tree —
+// see the layout diagram at the top of `crate::drive::shielded::paths`.
 
-/// The subtree key for per-block nullifiers storage as u8
-pub const SHIELDED_RECENT_NULLIFIERS_KEY_U8: u8 = b'n';
+/// The subtree key for per-block nullifiers storage (CountSumTree).
+///
+/// Depth 2 in the parent tree (right subtree of `SHIELDED_NOTES_KEY`).
+pub const SHIELDED_RECENT_NULLIFIERS_KEY: &[u8; 1] = &[160];
 
-/// The subtree key for compacted nullifiers storage
-pub const SHIELDED_COMPACTED_NULLIFIERS_KEY: &[u8; 1] = b"o";
+/// The subtree key for per-block nullifiers storage as u8.
+pub const SHIELDED_RECENT_NULLIFIERS_KEY_U8: u8 = 160;
 
-/// The subtree key for compacted nullifiers storage as u8
-pub const SHIELDED_COMPACTED_NULLIFIERS_KEY_U8: u8 = b'o';
+/// The subtree key for compacted nullifiers storage.
+///
+/// Depth 2 in the parent tree.
+pub const SHIELDED_COMPACTED_NULLIFIERS_KEY: &[u8; 1] = &[224];
 
-/// The subtree key for nullifiers expiration time storage
-pub const SHIELDED_NULLIFIERS_EXPIRATION_TIME_KEY: &[u8; 1] = b"p";
+/// The subtree key for compacted nullifiers storage as u8.
+pub const SHIELDED_COMPACTED_NULLIFIERS_KEY_U8: u8 = 224;
 
-/// The subtree key for nullifiers expiration time storage as u8
-pub const SHIELDED_NULLIFIERS_EXPIRATION_TIME_KEY_U8: u8 = b'p';
+/// The subtree key for nullifiers expiration time storage.
+///
+/// Deepest leaf in the parent tree — only touched by periodic expiry sweeps.
+pub const SHIELDED_NULLIFIERS_EXPIRATION_TIME_KEY: &[u8; 1] = &[240];
 
-/// Path to per-block nullifiers: [AddressBalances, "s", "n"]
+/// The subtree key for nullifiers expiration time storage as u8.
+pub const SHIELDED_NULLIFIERS_EXPIRATION_TIME_KEY_U8: u8 = 240;
+
+/// Path to per-block nullifiers: [AddressBalances, "s", [160]]
 pub fn shielded_recent_nullifiers_path() -> [&'static [u8]; 3] {
     [
         Into::<&[u8; 1]>::into(RootTree::AddressBalances),
@@ -28,7 +37,7 @@ pub fn shielded_recent_nullifiers_path() -> [&'static [u8]; 3] {
     ]
 }
 
-/// Path to per-block nullifiers as vec: [AddressBalances, "s", "n"]
+/// Path to per-block nullifiers as vec: [AddressBalances, "s", [160]]
 pub fn shielded_recent_nullifiers_path_vec() -> Vec<Vec<u8>> {
     vec![
         vec![RootTree::AddressBalances as u8],
@@ -37,7 +46,7 @@ pub fn shielded_recent_nullifiers_path_vec() -> Vec<Vec<u8>> {
     ]
 }
 
-/// Path to compacted nullifiers: [AddressBalances, "s", "o"]
+/// Path to compacted nullifiers: [AddressBalances, "s", [224]]
 pub fn shielded_compacted_nullifiers_path() -> [&'static [u8]; 3] {
     [
         Into::<&[u8; 1]>::into(RootTree::AddressBalances),
@@ -46,7 +55,7 @@ pub fn shielded_compacted_nullifiers_path() -> [&'static [u8]; 3] {
     ]
 }
 
-/// Path to compacted nullifiers as vec: [AddressBalances, "s", "o"]
+/// Path to compacted nullifiers as vec: [AddressBalances, "s", [224]]
 pub fn shielded_compacted_nullifiers_path_vec() -> Vec<Vec<u8>> {
     vec![
         vec![RootTree::AddressBalances as u8],
@@ -55,7 +64,7 @@ pub fn shielded_compacted_nullifiers_path_vec() -> Vec<Vec<u8>> {
     ]
 }
 
-/// Path to nullifiers expiration time: [AddressBalances, "s", "p"]
+/// Path to nullifiers expiration time: [AddressBalances, "s", [240]]
 pub fn shielded_nullifiers_expiration_time_path() -> [&'static [u8]; 3] {
     [
         Into::<&[u8; 1]>::into(RootTree::AddressBalances),
@@ -64,7 +73,7 @@ pub fn shielded_nullifiers_expiration_time_path() -> [&'static [u8]; 3] {
     ]
 }
 
-/// Path to nullifiers expiration time as vec: [AddressBalances, "s", "p"]
+/// Path to nullifiers expiration time as vec: [AddressBalances, "s", [240]]
 pub fn shielded_nullifiers_expiration_time_path_vec() -> Vec<Vec<u8>> {
     vec![
         vec![RootTree::AddressBalances as u8],