fix(drive-abci): record unshield and shield-surplus credits in recent address balance changes

packages/rs-drive-abci/src/abci/app/execution_result.rs

@@ -29,14 +29,14 @@ impl TryIntoPlatformVersioned<Option<ExecTxResult>> for StateTransitionExecution
                 gas_used: 0,
                 ..Default::default()
             }),
-            StateTransitionExecutionResult::PaidConsensusError { error, actual_fees } => {
-                Some(ExecTxResult {
-                    code: HandlerError::from(&error).code(),
-                    info: error.response_info_for_version(platform_version)?,
-                    gas_used: actual_fees.total_base_fee() as SignedCredits,
-                    ..Default::default()
-                })
-            }
+            StateTransitionExecutionResult::PaidConsensusError {
+                error, actual_fees, ..
+            } => Some(ExecTxResult {
+                code: HandlerError::from(&error).code(),
+                info: error.response_info_for_version(platform_version)?,
+                gas_used: actual_fees.total_base_fee() as SignedCredits,
+                ..Default::default()
+            }),
             StateTransitionExecutionResult::InternalError(message) => Some(ExecTxResult {
                 code: HandlerError::Internal(message).code(),
                 // TODO: That would be nice to provide more information about the error for debugging
@@ -118,6 +118,7 @@ mod tests {
         let execution_result = StateTransitionExecutionResult::PaidConsensusError {
             error: consensus_error,
             actual_fees: fee_result.clone(),
+            address_balance_changes: BTreeMap::new(),
         };
 
         let result: Option<ExecTxResult> = execution_result

packages/rs-drive-abci/src/execution/platform_events/state_transition_processing/execute_event/v0/mod.rs

@@ -1,5 +1,6 @@
 use crate::error::execution::ExecutionError;
 use crate::error::Error;
+use crate::execution::platform_events::state_transition_processing::record_added_balance_outputs::AddedBalanceOutputsOrigin;
 use crate::execution::types::execution_event::ExecutionEvent;
 use crate::execution::types::execution_operation::ValidationOperation;
 use crate::platform_types::event_execution_result::EventExecutionResult;
@@ -433,33 +434,15 @@ where
                     previous_fee_versions,
                 )?;
 
-                // Track address outputs if provided (e.g., for IdentityCreditTransferToAddresses)
-                if let Some(outputs) = added_to_balance_outputs {
-                    if let Some(balance_updates) = address_balances_in_update {
-                        for (address, credits) in outputs {
-                            let new_op = CreditOperation::AddToCredits(credits);
-                            balance_updates
-                                .entry(address)
-                                .and_modify(|existing| {
-                                    *existing = match existing {
-                                        // Set + Add = Set to combined value
-                                        CreditOperation::SetCredits(set_val) => {
-                                            CreditOperation::SetCredits(
-                                                set_val.saturating_add(credits),
-                                            )
-                                        }
-                                        // Add + Add = saturating add
-                                        CreditOperation::AddToCredits(add_val) => {
-                                            CreditOperation::AddToCredits(
-                                                add_val.saturating_add(credits),
-                                            )
-                                        }
-                                    };
-                                })
-                                .or_insert(new_op);
-                        }
-                    }
-                }
+                // Track address outputs if provided (e.g., for IdentityCreditTransferToAddresses).
+                // Transparent origin: this long-standing recording is preserved by every version
+                // (v0 and v1 both fold Transparent-origin outputs).
+                self.record_added_balance_outputs(
+                    address_balances_in_update,
+                    added_to_balance_outputs,
+                    AddedBalanceOutputsOrigin::Transparent,
+                    platform_version,
+                )?;
 
                 Ok(result)
             }
@@ -548,6 +531,7 @@ where
             ExecutionEvent::PaidFromShieldedPool {
                 operations,
                 fees_to_add_to_pool,
+                added_to_balance_outputs,
                 chargeable_failure,
             } => {
                 // An error-bearing `PaidFromShieldedPool` is legitimate ONLY for the
@@ -587,6 +571,18 @@ where
                     )
                     .map_err(Error::Drive)?;
 
+                // The ops just applied credited any transparent output address (an Unshield's
+                // recipient, including the chargeable-failure fallback address). Record that credit
+                // so incremental client sync sees it. ShieldedSpend origin: recorded only from v1
+                // (protocol v13); v0 drops it so pre-v13 blocks byte-match. Reached only on the
+                // applied path (the early return above skips this).
+                self.record_added_balance_outputs(
+                    address_balances_in_update,
+                    added_to_balance_outputs,
+                    AddedBalanceOutputsOrigin::ShieldedSpend,
+                    platform_version,
+                )?;
+
                 // Split the carved fee like every other transition: the real storage
                 // cost of the (permanent) shielded writes goes to the storage pool, so it
                 // is amortised to the validators that store it over time and picks up the
@@ -612,6 +608,7 @@ where
             }
             ExecutionEvent::PaidFromAssetLockToPool {
                 fees_to_add_to_pool,
+                added_to_balance_outputs,
                 operations,
                 ..
             } => {
@@ -633,6 +630,17 @@ where
                         )
                         .map_err(Error::Drive)?;
 
+                    // The ops just applied credited the shield's transparent surplus-output address
+                    // (when set). Record that credit so incremental client sync sees it. ShieldedSpend
+                    // origin: recorded only from v1 (protocol v13); v0 drops it so pre-v13 blocks
+                    // byte-match. Reached only on the applied (no-errors) path.
+                    self.record_added_balance_outputs(
+                        address_balances_in_update,
+                        added_to_balance_outputs,
+                        AddedBalanceOutputsOrigin::ShieldedSpend,
+                        platform_version,
+                    )?;
+
                     // Route the real storage cost of the shielded writes to the storage pool
                     // (amortised over time, epoch fee multiplier applied at payout); the
                     // remainder is the processing fee paid to the proposer. Conservation:
@@ -701,3 +709,147 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::platform_types::event_execution_result::EventExecutionResult;
+    use crate::test::helpers::setup::TestPlatformBuilder;
+    use dpp::fee::default_costs::CachedEpochIndexFeeVersions;
+
+    /// Regression for the recent-address-balance-changes gap: a `PaidFromShieldedPool` carrying an
+    /// Unshield's output credit must fold that credit into the `address_balances_in_update` map (the
+    /// sole feed for `store_address_balances_for_block`), so incremental client sync can see the
+    /// unshield. Before the fix the arm applied the drive ops but never touched the map, leaving the
+    /// credit invisible to incremental sync.
+    #[test]
+    fn paid_from_shielded_pool_folds_output_credit_into_address_balances() {
+        // Pin v13: the executor's record_added_balance_outputs is v1 here, which folds
+        // shielded-spend-origin outputs (v0 would drop them).
+        let platform_version = PlatformVersion::get(13).expect("v13 must exist");
+        let platform = TestPlatformBuilder::new()
+            .build_with_mock_rpc()
+            .set_genesis_state();
+        let transaction = platform.drive.grove.start_transaction();
+        let fee_versions = CachedEpochIndexFeeVersions::new();
+
+        let output_address = PlatformAddress::P2pkh([0xBB; 20]);
+        let event = ExecutionEvent::PaidFromShieldedPool {
+            operations: vec![],
+            fees_to_add_to_pool: 0,
+            added_to_balance_outputs: Some(BTreeMap::from([(output_address, 2500)])),
+            chargeable_failure: false,
+        };
+
+        let mut address_balances: BTreeMap<PlatformAddress, CreditOperation> = BTreeMap::new();
+        let result = platform
+            .platform
+            .execute_event_v0(
+                event,
+                vec![],
+                &BlockInfo::default(),
+                &transaction,
+                Some(&mut address_balances),
+                platform_version,
+                &fee_versions,
+            )
+            .expect("execute_event_v0 should not error");
+
+        assert!(
+            matches!(result, EventExecutionResult::SuccessfulPaidExecution(..)),
+            "an ordinary unshield must execute successfully"
+        );
+        assert_eq!(
+            address_balances.get(&output_address),
+            Some(&CreditOperation::AddToCredits(2500)),
+            "the unshield output credit must be recorded for incremental sync"
+        );
+    }
+
+    /// A `PaidFromShieldedPool` with no output (ShieldedTransfer / ShieldedWithdrawal, or a net-zero
+    /// unshield) must leave the address-balances map untouched.
+    #[test]
+    fn paid_from_shielded_pool_without_output_records_nothing() {
+        // Pin v13: the executor's record_added_balance_outputs is v1 here, which folds
+        // shielded-spend-origin outputs (v0 would drop them).
+        let platform_version = PlatformVersion::get(13).expect("v13 must exist");
+        let platform = TestPlatformBuilder::new()
+            .build_with_mock_rpc()
+            .set_genesis_state();
+        let transaction = platform.drive.grove.start_transaction();
+        let fee_versions = CachedEpochIndexFeeVersions::new();
+
+        let event = ExecutionEvent::PaidFromShieldedPool {
+            operations: vec![],
+            fees_to_add_to_pool: 0,
+            added_to_balance_outputs: None,
+            chargeable_failure: false,
+        };
+
+        let mut address_balances: BTreeMap<PlatformAddress, CreditOperation> = BTreeMap::new();
+        platform
+            .platform
+            .execute_event_v0(
+                event,
+                vec![],
+                &BlockInfo::default(),
+                &transaction,
+                Some(&mut address_balances),
+                platform_version,
+                &fee_versions,
+            )
+            .expect("execute_event_v0 should not error");
+
+        assert!(
+            address_balances.is_empty(),
+            "a shielded spend crediting no transparent address must record nothing"
+        );
+    }
+
+    /// The `PaidFromAssetLockToPool` counterpart: a `ShieldFromAssetLock` surplus credit must be
+    /// folded into the address-balances map on the applied (no-errors) path.
+    #[test]
+    fn paid_from_asset_lock_to_pool_folds_surplus_credit_into_address_balances() {
+        // Pin v13: the executor's record_added_balance_outputs is v1 here, which folds
+        // shielded-spend-origin outputs (v0 would drop them).
+        let platform_version = PlatformVersion::get(13).expect("v13 must exist");
+        let platform = TestPlatformBuilder::new()
+            .with_latest_protocol_version()
+            .build_with_mock_rpc()
+            .set_initial_state_structure();
+        let transaction = platform.drive.grove.start_transaction();
+        let fee_versions = CachedEpochIndexFeeVersions::new();
+
+        let surplus_address = PlatformAddress::P2pkh([0x42; 20]);
+        let event = ExecutionEvent::PaidFromAssetLockToPool {
+            fees_to_add_to_pool: 1_000_000,
+            added_to_balance_outputs: Some(BTreeMap::from([(surplus_address, 2000)])),
+            operations: vec![],
+            execution_operations: vec![],
+        };
+
+        let mut address_balances: BTreeMap<PlatformAddress, CreditOperation> = BTreeMap::new();
+        let result = platform
+            .platform
+            .execute_event_v0(
+                event,
+                vec![],
+                &BlockInfo::default(),
+                &transaction,
+                Some(&mut address_balances),
+                platform_version,
+                &fee_versions,
+            )
+            .expect("execute_event_v0 should not error");
+
+        assert!(
+            matches!(result, EventExecutionResult::SuccessfulPaidExecution(..)),
+            "a sufficiently-funded shield-from-asset-lock must execute successfully"
+        );
+        assert_eq!(
+            address_balances.get(&surplus_address),
+            Some(&CreditOperation::AddToCredits(2000)),
+            "the shield surplus credit must be recorded for incremental sync"
+        );
+    }
+}

packages/rs-drive-abci/src/execution/platform_events/state_transition_processing/mod.rs

@@ -2,5 +2,6 @@ mod cleanup_recent_block_storage_address_balances;
 mod decode_raw_state_transitions;
 mod execute_event;
 mod process_raw_state_transitions;
+mod record_added_balance_outputs;
 mod store_address_balances_to_recent_block_storage;
 mod validate_fees_of_event;