sui_json_rpc/

authority_state.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use arc_swap::Guard;
use async_trait::async_trait;
use move_core_types::language_storage::TypeTag;
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use sui_core::accumulators::balances::{get_all_balances_for_owner, get_balance};
use sui_core::authority::AuthorityState;
use sui_core::authority::authority_per_epoch_store::AuthorityPerEpochStore;
use sui_core::execution_cache::ObjectCacheRead;
use sui_core::jsonrpc_index::{CoinIndexKey2, CoinInfo, TotalBalance};
use sui_core::subscription_handler::SubscriptionHandler;
use sui_json_rpc_types::{
    Coin as SuiCoin, DevInspectResults, DryRunTransactionBlockResponse, EventFilter, SuiEvent,
    SuiObjectDataFilter, TransactionFilter,
};
use sui_storage::key_value_store::{
    KVStoreTransactionData, TransactionKeyValueStore, TransactionKeyValueStoreTrait,
};
use sui_types::accumulator_root::AccumulatorKey;
use sui_types::balance::Balance;
use sui_types::base_types::{
    MoveObjectType, ObjectID, ObjectInfo, ObjectRef, SequenceNumber, SuiAddress,
};
use sui_types::bridge::Bridge;
use sui_types::coin_reservation;
use sui_types::committee::{Committee, EpochId};
use sui_types::digests::{ChainIdentifier, TransactionDigest};
use sui_types::dynamic_field::DynamicFieldInfo;
use sui_types::effects::TransactionEffects;
use sui_types::error::{SuiError, SuiErrorKind, SuiResult, UserInputError};
use sui_types::event::EventID;
use sui_types::governance::StakedSui;
use sui_types::messages_checkpoint::{
    CheckpointContents, CheckpointContentsDigest, CheckpointDigest, CheckpointSequenceNumber,
    VerifiedCheckpoint,
};
use sui_types::object::{MoveObject, Object, ObjectRead, Owner, PastObjectRead};
use sui_types::storage::{BackingPackageStore, ObjectStore, WriteKind};
use sui_types::sui_serde::BigInt;
use sui_types::sui_system_state::SuiSystemState;
use sui_types::transaction::{Transaction, TransactionData, TransactionKind};
use thiserror::Error;
use tokio::task::JoinError;

use crate::ObjectProvider;
#[cfg(test)]
use mockall::automock;
use typed_store_error::TypedStoreError;

pub type StateReadResult<T = ()> = Result<T, StateReadError>;

/// Trait for AuthorityState methods commonly used by at least two api.
#[cfg_attr(test, automock)]
#[async_trait]
pub trait StateRead: Send + Sync {
    async fn multi_get(
        &self,
        transactions: &[TransactionDigest],
        effects: &[TransactionDigest],
    ) -> StateReadResult<KVStoreTransactionData>;

    fn get_object_read(&self, object_id: &ObjectID) -> StateReadResult<ObjectRead>;

    fn get_past_object_read(
        &self,
        object_id: &ObjectID,
        version: SequenceNumber,
    ) -> StateReadResult<PastObjectRead>;

    async fn get_object(&self, object_id: &ObjectID) -> StateReadResult<Option<Object>>;

    fn load_epoch_store_one_call_per_task(&self) -> Guard<Arc<AuthorityPerEpochStore>>;

    fn get_dynamic_fields(
        &self,
        owner: ObjectID,
        cursor: Option<ObjectID>,
        limit: usize,
    ) -> StateReadResult<Vec<(ObjectID, DynamicFieldInfo)>>;

    fn get_cache_reader(&self) -> &Arc<dyn ObjectCacheRead>;

    fn get_object_store(&self) -> &Arc<dyn ObjectStore + Send + Sync>;

    fn get_backing_package_store(&self) -> &Arc<dyn BackingPackageStore + Send + Sync>;

    fn get_owner_objects(
        &self,
        owner: SuiAddress,
        cursor: Option<ObjectID>,
        filter: Option<SuiObjectDataFilter>,
    ) -> StateReadResult<Vec<ObjectInfo>>;

    async fn query_events(
        &self,
        kv_store: &Arc<TransactionKeyValueStore>,
        query: EventFilter,
        // If `Some`, the query will start from the next item after the specified cursor
        cursor: Option<EventID>,
        limit: usize,
        descending: bool,
    ) -> StateReadResult<Vec<SuiEvent>>;

    // transaction_execution_api
    #[allow(clippy::type_complexity)]
    async fn dry_exec_transaction(
        &self,
        transaction: TransactionData,
        transaction_digest: TransactionDigest,
    ) -> StateReadResult<(
        DryRunTransactionBlockResponse,
        BTreeMap<ObjectID, (ObjectRef, Object, WriteKind)>,
        TransactionEffects,
        Option<ObjectID>,
    )>;

    async fn dev_inspect_transaction_block(
        &self,
        sender: SuiAddress,
        transaction_kind: TransactionKind,
        gas_price: Option<u64>,
        gas_budget: Option<u64>,
        gas_sponsor: Option<SuiAddress>,
        gas_objects: Option<Vec<ObjectRef>>,
        show_raw_txn_data_and_effects: Option<bool>,
        skip_checks: Option<bool>,
    ) -> StateReadResult<DevInspectResults>;

    // indexer_api
    fn get_subscription_handler(&self) -> Arc<SubscriptionHandler>;

    fn get_owner_objects_with_limit(
        &self,
        owner: SuiAddress,
        cursor: Option<ObjectID>,
        limit: usize,
        filter: Option<SuiObjectDataFilter>,
    ) -> StateReadResult<Vec<ObjectInfo>>;

    async fn get_transactions(
        &self,
        kv_store: &Arc<TransactionKeyValueStore>,
        filter: Option<TransactionFilter>,
        cursor: Option<TransactionDigest>,
        limit: Option<usize>,
        reverse: bool,
    ) -> StateReadResult<Vec<TransactionDigest>>;

    fn get_dynamic_field_object_id(
        &self,
        owner: ObjectID,
        name_type: TypeTag,
        name_bcs_bytes: &[u8],
    ) -> StateReadResult<Option<ObjectID>>;

    // governance_api
    async fn get_staked_sui(&self, owner: SuiAddress) -> StateReadResult<Vec<StakedSui>>;
    fn get_system_state(&self) -> StateReadResult<SuiSystemState>;
    fn get_or_latest_committee(&self, epoch: Option<BigInt<u64>>) -> StateReadResult<Committee>;

    // bridge_api
    fn get_bridge(&self) -> StateReadResult<Bridge>;

    // coin_api
    fn find_publish_txn_digest(&self, package_id: ObjectID) -> StateReadResult<TransactionDigest>;
    fn get_owned_coins(
        &self,
        owner: SuiAddress,
        cursor: (String, u64, ObjectID),
        limit: usize,
        one_coin_type_only: bool,
    ) -> StateReadResult<Vec<SuiCoin>>;
    async fn get_executed_transaction_and_effects(
        &self,
        digest: TransactionDigest,
        kv_store: Arc<TransactionKeyValueStore>,
    ) -> StateReadResult<(Transaction, TransactionEffects)>;
    async fn get_balance(
        &self,
        owner: SuiAddress,
        coin_type: TypeTag,
    ) -> StateReadResult<TotalBalance>;
    async fn get_all_balance(
        &self,
        owner: SuiAddress,
    ) -> StateReadResult<Arc<HashMap<TypeTag, TotalBalance>>>;

    // read_api
    fn get_verified_checkpoint_by_sequence_number(
        &self,
        sequence_number: CheckpointSequenceNumber,
    ) -> StateReadResult<VerifiedCheckpoint>;

    fn get_checkpoint_contents(
        &self,
        digest: CheckpointContentsDigest,
    ) -> StateReadResult<CheckpointContents>;

    fn get_verified_checkpoint_summary_by_digest(
        &self,
        digest: CheckpointDigest,
    ) -> StateReadResult<VerifiedCheckpoint>;

    fn deprecated_multi_get_transaction_checkpoint(
        &self,
        digests: &[TransactionDigest],
    ) -> StateReadResult<Vec<Option<(EpochId, CheckpointSequenceNumber)>>>;

    fn deprecated_get_transaction_checkpoint(
        &self,
        digest: &TransactionDigest,
    ) -> StateReadResult<Option<(EpochId, CheckpointSequenceNumber)>>;

    fn multi_get_checkpoint_by_sequence_number(
        &self,
        sequence_numbers: &[CheckpointSequenceNumber],
    ) -> StateReadResult<Vec<Option<VerifiedCheckpoint>>>;

    fn get_total_transaction_blocks(&self) -> StateReadResult<u64>;

    fn get_checkpoint_by_sequence_number(
        &self,
        sequence_number: CheckpointSequenceNumber,
    ) -> StateReadResult<Option<VerifiedCheckpoint>>;

    fn get_latest_checkpoint_sequence_number(&self) -> StateReadResult<CheckpointSequenceNumber>;

    fn get_chain_identifier(&self) -> StateReadResult<ChainIdentifier>;
}

#[async_trait]
impl StateRead for AuthorityState {
    async fn multi_get(
        &self,
        transactions: &[TransactionDigest],
        effects: &[TransactionDigest],
    ) -> StateReadResult<KVStoreTransactionData> {
        Ok(
            <AuthorityState as TransactionKeyValueStoreTrait>::multi_get(
                self,
                transactions,
                effects,
            )
            .await?,
        )
    }

    fn get_object_read(&self, object_id: &ObjectID) -> StateReadResult<ObjectRead> {
        let result = self.get_object_read(object_id)?;

        // If object not found and coin reservations are enabled, check if this is a
        // masked object ID (fake coin request).
        if let ObjectRead::NotExists(object_id) = result
            && self
                .load_epoch_store_one_call_per_task()
                .protocol_config()
                .enable_coin_reservation_obj_refs()
        {
            let chain_identifier = self.get_chain_identifier();
            let unmasked_id = coin_reservation::mask_or_unmask_id(object_id, chain_identifier);

            // Try to load the unmasked object (the accumulator)
            if let ObjectRead::Exists(_, object, _) = self.get_object_read(&unmasked_id)? {
                let accumulator_version = object.version();
                let Some(move_object) = object.data.try_as_move() else {
                    // Not a move object, return original NotExists
                    return Ok(ObjectRead::NotExists(object_id));
                };
                let Some(currency_type) =
                    move_object.type_().balance_accumulator_field_type_maybe()
                else {
                    // Not an accumulator object, return original NotExists
                    return Ok(ObjectRead::NotExists(object_id));
                };

                let balance_type = Balance::type_tag(currency_type.clone());

                let (AccumulatorKey { owner }, value) = move_object.try_into()?;

                let Some((object_ref, balance, previous_transaction)) =
                    self.get_address_balance_coin_info(owner, balance_type)?
                else {
                    return Ok(ObjectRead::NotExists(object_id));
                };

                debug_assert_eq!(balance, value.as_u128().map(|v| v as u64).unwrap_or(0));

                // Create a fake coin object with the masked ID
                let coin = Object::new_move(
                    MoveObject::new_coin(currency_type, accumulator_version, object_id, balance),
                    Owner::AddressOwner(owner),
                    previous_transaction,
                );

                let layout = self.get_object_layout(&coin)?;
                return Ok(ObjectRead::Exists(object_ref, coin, layout));
            }

            return Ok(ObjectRead::NotExists(object_id));
        }

        Ok(result)
    }

    async fn get_object(&self, object_id: &ObjectID) -> StateReadResult<Option<Object>> {
        Ok(self.get_object(object_id).await)
    }

    fn get_past_object_read(
        &self,
        object_id: &ObjectID,
        version: SequenceNumber,
    ) -> StateReadResult<PastObjectRead> {
        Ok(self.get_past_object_read(object_id, version)?)
    }

    fn load_epoch_store_one_call_per_task(&self) -> Guard<Arc<AuthorityPerEpochStore>> {
        self.load_epoch_store_one_call_per_task()
    }

    fn get_dynamic_fields(
        &self,
        owner: ObjectID,
        cursor: Option<ObjectID>,
        limit: usize,
    ) -> StateReadResult<Vec<(ObjectID, DynamicFieldInfo)>> {
        Ok(self.get_dynamic_fields(owner, cursor, limit)?)
    }

    fn get_cache_reader(&self) -> &Arc<dyn ObjectCacheRead> {
        self.get_object_cache_reader()
    }

    fn get_object_store(&self) -> &Arc<dyn ObjectStore + Send + Sync> {
        self.get_object_store()
    }

    fn get_backing_package_store(&self) -> &Arc<dyn BackingPackageStore + Send + Sync> {
        self.get_backing_package_store()
    }

    fn get_owner_objects(
        &self,
        owner: SuiAddress,
        cursor: Option<ObjectID>,
        filter: Option<SuiObjectDataFilter>,
    ) -> StateReadResult<Vec<ObjectInfo>> {
        Ok(self
            .get_owner_objects_iterator(owner, cursor, filter)?
            .collect())
    }

    async fn query_events(
        &self,
        kv_store: &Arc<TransactionKeyValueStore>,
        query: EventFilter,
        // If `Some`, the query will start from the next item after the specified cursor
        cursor: Option<EventID>,
        limit: usize,
        descending: bool,
    ) -> StateReadResult<Vec<SuiEvent>> {
        Ok(self
            .query_events(kv_store, query, cursor, limit, descending)
            .await?)
    }

    #[allow(clippy::type_complexity)]
    async fn dry_exec_transaction(
        &self,
        transaction: TransactionData,
        transaction_digest: TransactionDigest,
    ) -> StateReadResult<(
        DryRunTransactionBlockResponse,
        BTreeMap<ObjectID, (ObjectRef, Object, WriteKind)>,
        TransactionEffects,
        Option<ObjectID>,
    )> {
        Ok(self
            .dry_exec_transaction(transaction, transaction_digest)
            .await?)
    }

    async fn dev_inspect_transaction_block(
        &self,
        sender: SuiAddress,
        transaction_kind: TransactionKind,
        gas_price: Option<u64>,
        gas_budget: Option<u64>,
        gas_sponsor: Option<SuiAddress>,
        gas_objects: Option<Vec<ObjectRef>>,
        show_raw_txn_data_and_effects: Option<bool>,
        skip_checks: Option<bool>,
    ) -> StateReadResult<DevInspectResults> {
        Ok(self
            .dev_inspect_transaction_block(
                sender,
                transaction_kind,
                gas_price,
                gas_budget,
                gas_sponsor,
                gas_objects,
                show_raw_txn_data_and_effects,
                skip_checks,
            )
            .await?)
    }

    fn get_subscription_handler(&self) -> Arc<SubscriptionHandler> {
        self.subscription_handler.clone()
    }

    fn get_owner_objects_with_limit(
        &self,
        owner: SuiAddress,
        cursor: Option<ObjectID>,
        limit: usize,
        filter: Option<SuiObjectDataFilter>,
    ) -> StateReadResult<Vec<ObjectInfo>> {
        Ok(self.get_owner_objects(owner, cursor, limit, filter)?)
    }

    async fn get_transactions(
        &self,
        kv_store: &Arc<TransactionKeyValueStore>,
        filter: Option<TransactionFilter>,
        cursor: Option<TransactionDigest>,
        limit: Option<usize>,
        reverse: bool,
    ) -> StateReadResult<Vec<TransactionDigest>> {
        Ok(self
            .get_transactions(kv_store, filter, cursor, limit, reverse)
            .await?)
    }

    fn get_dynamic_field_object_id(
        // indexer
        &self,
        owner: ObjectID,
        name_type: TypeTag,
        name_bcs_bytes: &[u8],
    ) -> StateReadResult<Option<ObjectID>> {
        Ok(self.get_dynamic_field_object_id(owner, name_type, name_bcs_bytes)?)
    }

    async fn get_staked_sui(&self, owner: SuiAddress) -> StateReadResult<Vec<StakedSui>> {
        Ok(self
            .get_move_objects(owner, MoveObjectType::staked_sui())
            .await?)
    }
    fn get_system_state(&self) -> StateReadResult<SuiSystemState> {
        Ok(self
            .get_object_cache_reader()
            .get_sui_system_state_object_unsafe()?)
    }
    fn get_or_latest_committee(&self, epoch: Option<BigInt<u64>>) -> StateReadResult<Committee> {
        Ok(self
            .committee_store()
            .get_or_latest_committee(epoch.map(|e| *e))?)
    }

    fn get_bridge(&self) -> StateReadResult<Bridge> {
        self.get_cache_reader()
            .get_bridge_object_unsafe()
            .map_err(|err| err.into())
    }

    fn find_publish_txn_digest(&self, package_id: ObjectID) -> StateReadResult<TransactionDigest> {
        Ok(self.find_publish_txn_digest(package_id)?)
    }
    fn get_owned_coins(
        &self,
        owner: SuiAddress,
        cursor: (String, u64, ObjectID),
        limit: usize,
        one_coin_type_only: bool,
    ) -> StateReadResult<Vec<SuiCoin>> {
        // Ordering per coin type: [real[0], fake, real[1], real[2], ...]
        // The fake coin (address balance) is always at position 1 within its type.

        fn to_sui_coin(key: CoinIndexKey2, info: CoinInfo) -> SuiCoin {
            SuiCoin {
                coin_type: key.coin_type,
                coin_object_id: key.object_id,
                version: info.version,
                digest: info.digest,
                balance: info.balance,
                previous_transaction: info.previous_transaction,
            }
        }

        fn obj_ref_to_sui_coin(
            coin_type: String,
            obj_ref: ObjectRef,
            balance: u64,
            previous_transaction: TransactionDigest,
        ) -> SuiCoin {
            SuiCoin {
                coin_type,
                coin_object_id: obj_ref.0,
                version: obj_ref.1,
                digest: obj_ref.2,
                balance,
                previous_transaction,
            }
        }

        // Build fake coins map (only when coin reservations are enabled).
        let coin_reservations_enabled = self
            .load_epoch_store_one_call_per_task()
            .protocol_config()
            .enable_coin_reservation_obj_refs();

        let fake_coins: HashMap<String, SuiCoin> = if !coin_reservations_enabled {
            HashMap::new()
        } else if one_coin_type_only {
            let balance_type_tag = sui_types::parse_sui_type_tag(&cursor.0)
                .map_err(|e| anyhow::anyhow!("Invalid coin type: {} - {}", cursor.0, e))?;
            let balance_type = Balance::type_tag(balance_type_tag);
            self.get_address_balance_coin_info(owner, balance_type)?
                .map(|(obj_ref, balance, prev_tx)| {
                    HashMap::from([(
                        cursor.0.clone(),
                        obj_ref_to_sui_coin(cursor.0.clone(), obj_ref, balance, prev_tx),
                    )])
                })
                .unwrap_or_default()
        } else {
            self.get_all_address_balance_coin_infos(owner)?
                .into_iter()
                .map(|(coin_type, (obj_ref, balance, prev_tx))| {
                    (
                        coin_type.clone(),
                        obj_ref_to_sui_coin(coin_type, obj_ref, balance, prev_tx),
                    )
                })
                .collect()
        };

        // Determine cursor state.
        let cursor_at_fake = fake_coins.values().any(|c| c.coin_object_id == cursor.2);

        // If cursor is at fake coin, reset to start of that type and skip real[0].
        let (real_cursor, skip_first_real) = if cursor_at_fake {
            ((cursor.0.clone(), 0, ObjectID::ZERO), true)
        } else {
            (cursor.clone(), false)
        };

        let real_coins_iter = self.get_owned_coins_iterator_with_cursor(
            owner,
            real_cursor.clone(),
            limit + 1,
            one_coin_type_only,
        )?;

        // Track which types have had their fake coin emitted.
        let mut fake_emitted: HashMap<String, bool> = HashMap::new();

        // If cursor is a real coin, check if we're past the fake coin slot.
        // The fake coin is at position 1 (after first real). So if cursor is the
        // first real coin, fake hasn't been emitted. If cursor is any later real
        // coin, fake was already emitted.
        let mut emit_fake_before_reals = false;
        if cursor.2 != ObjectID::ZERO && !cursor_at_fake && fake_coins.contains_key(&cursor.0) {
            // Check if cursor is the first real coin by querying from the start
            let first_real_id = self
                .get_owned_coins_iterator_with_cursor(
                    owner,
                    (cursor.0.clone(), 0, ObjectID::ZERO),
                    1,
                    one_coin_type_only,
                )?
                .next()
                .map(|(k, _)| k.object_id);

            if first_real_id == Some(cursor.2) {
                // Cursor is at first real coin, fake should be emitted next
                emit_fake_before_reals = true;
            } else {
                // Cursor is past first real coin, fake was already emitted
                fake_emitted.insert(cursor.0.clone(), true);
            }
        }

        let mut result = Vec::with_capacity(limit);

        // If cursor is at first real coin, emit fake before continuing with more reals
        if emit_fake_before_reals && let Some(fake) = fake_coins.get(&cursor.0) {
            result.push(fake.clone());
            fake_emitted.insert(cursor.0.clone(), true);
        }

        let mut seen_first_real: HashMap<String, bool> = HashMap::new();
        let mut skipped_first = false;

        for (key, info) in real_coins_iter {
            if result.len() >= limit {
                break;
            }

            let coin = to_sui_coin(key, info);
            let coin_type = &coin.coin_type;
            let is_first_real = !seen_first_real.get(coin_type).copied().unwrap_or(false);

            // Skip first real coin when resuming from a fake coin cursor.
            if skip_first_real && !skipped_first && coin_type == &real_cursor.0 {
                skipped_first = true;
                seen_first_real.insert(coin_type.clone(), true);
                continue;
            }

            // Emit the real coin.
            result.push(coin.clone());
            seen_first_real.insert(coin_type.clone(), true);

            // After first real coin of a type, emit its fake coin (if not already emitted).
            if is_first_real && !fake_emitted.get(coin_type).copied().unwrap_or(false) {
                if let Some(fake) = fake_coins.get(coin_type)
                    && result.len() < limit
                {
                    result.push(fake.clone());
                }
                fake_emitted.insert(coin_type.clone(), true);
            }
        }

        // Emit any fake coins for types that had no real coins.
        // Only do this on the first page (cursor at start) to avoid duplicates.
        if cursor.2 == ObjectID::ZERO {
            for (coin_type, fake) in fake_coins {
                if result.len() >= limit {
                    break;
                }
                if !fake_emitted.get(&coin_type).copied().unwrap_or(false) {
                    result.push(fake);
                }
            }
        }

        Ok(result)
    }

    async fn get_executed_transaction_and_effects(
        &self,
        digest: TransactionDigest,
        kv_store: Arc<TransactionKeyValueStore>,
    ) -> StateReadResult<(Transaction, TransactionEffects)> {
        Ok(self
            .get_executed_transaction_and_effects(digest, kv_store)
            .await?)
    }

    async fn get_balance(
        &self,
        owner: SuiAddress,
        coin_type: TypeTag,
    ) -> StateReadResult<TotalBalance> {
        let indexes = self.indexes.clone();
        let child_object_resolver = self.get_child_object_resolver().clone();
        Ok(
            tokio::task::spawn_blocking(move || -> SuiResult<TotalBalance> {
                let address_balance =
                    get_balance(owner, child_object_resolver.as_ref(), coin_type.clone())?;
                let coin_balance = indexes
                    .as_ref()
                    .ok_or(SuiErrorKind::IndexStoreNotAvailable)?
                    .get_coin_object_balance(owner, coin_type)?;
                let mut total_balance = coin_balance;
                if address_balance > 0 {
                    total_balance.balance += address_balance as i128;
                    total_balance.num_coins += 1;
                }
                total_balance.address_balance = address_balance;
                Ok(total_balance)
            })
            .await
            .map_err(|e: JoinError| {
                SuiError(Box::new(SuiErrorKind::ExecutionError(e.to_string())))
            })??,
        )
    }

    async fn get_all_balance(
        &self,
        owner: SuiAddress,
    ) -> StateReadResult<Arc<HashMap<TypeTag, TotalBalance>>> {
        let indexes = self.indexes.clone();
        let child_object_resolver = self.get_child_object_resolver().clone();
        Ok(tokio::task::spawn_blocking(
            move || -> SuiResult<Arc<HashMap<TypeTag, TotalBalance>>> {
                let indexes = indexes
                    .as_ref()
                    .ok_or(SuiErrorKind::IndexStoreNotAvailable)?;
                let address_balances =
                    get_all_balances_for_owner(owner, child_object_resolver.as_ref(), indexes)?;
                let coin_balances = (*indexes.get_all_coin_object_balances(owner)?).clone();
                let mut all_balances = coin_balances;
                for (coin_type, balance) in address_balances {
                    let existing_balance = all_balances.entry(coin_type).or_insert(TotalBalance {
                        balance: 0,
                        num_coins: 0,
                        address_balance: 0,
                    });
                    existing_balance.balance += balance as i128;
                    existing_balance.num_coins += 1;
                    existing_balance.address_balance = balance;
                }
                Ok(Arc::new(all_balances))
            },
        )
        .await
        .map_err(|e: JoinError| {
            SuiError(Box::new(SuiErrorKind::ExecutionError(e.to_string())))
        })??)
    }

    fn get_verified_checkpoint_by_sequence_number(
        &self,
        sequence_number: CheckpointSequenceNumber,
    ) -> StateReadResult<VerifiedCheckpoint> {
        Ok(self.get_verified_checkpoint_by_sequence_number(sequence_number)?)
    }

    fn get_checkpoint_contents(
        &self,
        digest: CheckpointContentsDigest,
    ) -> StateReadResult<CheckpointContents> {
        Ok(self.get_checkpoint_contents(digest)?)
    }

    fn get_verified_checkpoint_summary_by_digest(
        &self,
        digest: CheckpointDigest,
    ) -> StateReadResult<VerifiedCheckpoint> {
        Ok(self.get_verified_checkpoint_summary_by_digest(digest)?)
    }

    fn deprecated_multi_get_transaction_checkpoint(
        &self,
        digests: &[TransactionDigest],
    ) -> StateReadResult<Vec<Option<(EpochId, CheckpointSequenceNumber)>>> {
        Ok(self
            .get_checkpoint_cache()
            .deprecated_multi_get_transaction_checkpoint(digests))
    }

    fn deprecated_get_transaction_checkpoint(
        &self,
        digest: &TransactionDigest,
    ) -> StateReadResult<Option<(EpochId, CheckpointSequenceNumber)>> {
        Ok(self
            .get_checkpoint_cache()
            .deprecated_get_transaction_checkpoint(digest))
    }

    fn multi_get_checkpoint_by_sequence_number(
        &self,
        sequence_numbers: &[CheckpointSequenceNumber],
    ) -> StateReadResult<Vec<Option<VerifiedCheckpoint>>> {
        Ok(self.multi_get_checkpoint_by_sequence_number(sequence_numbers)?)
    }

    fn get_total_transaction_blocks(&self) -> StateReadResult<u64> {
        Ok(self.get_total_transaction_blocks()?)
    }

    fn get_checkpoint_by_sequence_number(
        &self,
        sequence_number: CheckpointSequenceNumber,
    ) -> StateReadResult<Option<VerifiedCheckpoint>> {
        Ok(self.get_checkpoint_by_sequence_number(sequence_number)?)
    }

    fn get_latest_checkpoint_sequence_number(&self) -> StateReadResult<CheckpointSequenceNumber> {
        Ok(self.get_latest_checkpoint_sequence_number()?)
    }

    fn get_chain_identifier(&self) -> StateReadResult<ChainIdentifier> {
        Ok(self.get_chain_identifier())
    }
}

/// This implementation allows `S` to be a dynamically sized type (DST) that implements ObjectProvider
/// Valid as `S` is referenced only, and memory management is handled by `Arc`
#[async_trait]
impl<S: ?Sized + StateRead> ObjectProvider for Arc<S> {
    type Error = StateReadError;

    async fn get_object(
        &self,
        id: &ObjectID,
        version: &SequenceNumber,
    ) -> Result<Object, Self::Error> {
        Ok(self.get_past_object_read(id, *version)?.into_object()?)
    }

    async fn find_object_lt_or_eq_version(
        &self,
        id: &ObjectID,
        version: &SequenceNumber,
    ) -> Result<Option<Object>, Self::Error> {
        Ok(self
            .get_cache_reader()
            .find_object_lt_or_eq_version(*id, *version))
    }
}

#[async_trait]
impl<S: ?Sized + StateRead> ObjectProvider for (Arc<S>, Arc<TransactionKeyValueStore>) {
    type Error = StateReadError;

    async fn get_object(
        &self,
        id: &ObjectID,
        version: &SequenceNumber,
    ) -> Result<Object, Self::Error> {
        let object_read = self.0.get_past_object_read(id, *version)?;
        match object_read {
            PastObjectRead::ObjectNotExists(_) | PastObjectRead::VersionNotFound(..) => {
                match self.1.get_object(*id, *version).await? {
                    Some(object) => Ok(object),
                    None => Ok(PastObjectRead::VersionNotFound(*id, *version).into_object()?),
                }
            }
            _ => Ok(object_read.into_object()?),
        }
    }

    async fn find_object_lt_or_eq_version(
        &self,
        id: &ObjectID,
        version: &SequenceNumber,
    ) -> Result<Option<Object>, Self::Error> {
        Ok(self
            .0
            .get_cache_reader()
            .find_object_lt_or_eq_version(*id, *version))
    }
}

#[derive(Debug, Error)]
pub enum StateReadInternalError {
    #[error(transparent)]
    SuiError(#[from] SuiError),
    #[error(transparent)]
    JoinError(#[from] JoinError),
    #[error(transparent)]
    Anyhow(#[from] anyhow::Error),
}

impl From<SuiErrorKind> for StateReadInternalError {
    fn from(e: SuiErrorKind) -> Self {
        StateReadInternalError::SuiError(SuiError::from(e))
    }
}

#[derive(Debug, Error)]
pub enum StateReadClientError {
    #[error(transparent)]
    SuiError(#[from] SuiError),
    #[error(transparent)]
    UserInputError(#[from] UserInputError),
}

impl From<SuiErrorKind> for StateReadClientError {
    fn from(e: SuiErrorKind) -> Self {
        StateReadClientError::SuiError(SuiError::from(e))
    }
}

/// `StateReadError` is the error type for callers to work with.
/// It captures all possible errors that can occur while reading state, classifying them into two categories.
/// Unless `StateReadError` is the final error state before returning to caller, the app may still want error context.
/// This context is preserved in `Internal` and `Client` variants.
#[derive(Debug, Error)]
pub enum StateReadError {
    // sui_json_rpc::Error will do the final conversion to generic error message
    #[error(transparent)]
    Internal(#[from] StateReadInternalError),

    // Client errors
    #[error(transparent)]
    Client(#[from] StateReadClientError),
}

impl From<SuiErrorKind> for StateReadError {
    fn from(e: SuiErrorKind) -> Self {
        match e {
            SuiErrorKind::IndexStoreNotAvailable
            | SuiErrorKind::TransactionNotFound { .. }
            | SuiErrorKind::UnsupportedFeatureError { .. }
            | SuiErrorKind::UserInputError { .. }
            | SuiErrorKind::WrongMessageVersion { .. } => StateReadError::Client(e.into()),
            _ => StateReadError::Internal(e.into()),
        }
    }
}

impl From<SuiError> for StateReadError {
    fn from(e: SuiError) -> Self {
        e.into_inner().into()
    }
}

impl From<UserInputError> for StateReadError {
    fn from(e: UserInputError) -> Self {
        StateReadError::Client(e.into())
    }
}

impl From<JoinError> for StateReadError {
    fn from(e: JoinError) -> Self {
        StateReadError::Internal(e.into())
    }
}

impl From<anyhow::Error> for StateReadError {
    fn from(e: anyhow::Error) -> Self {
        StateReadError::Internal(e.into())
    }
}

impl From<TypedStoreError> for StateReadError {
    fn from(e: TypedStoreError) -> Self {
        let error: SuiError = e.into();
        StateReadError::Internal(error.into())
    }
}