sui_core/accumulators/
balances.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use mysten_common::debug_fatal;
5use sui_types::{
6    TypeTag,
7    accumulator_metadata::AccumulatorOwner,
8    accumulator_root::AccumulatorValue,
9    balance::Balance,
10    base_types::{ObjectID, SuiAddress},
11    error::SuiResult,
12    storage::error::Error as StorageError,
13    storage::{ChildObjectResolver, DynamicFieldIteratorItem, DynamicFieldKey, RpcIndexes},
14};
15
16use crate::jsonrpc_index::IndexStoreTables;
17
18/// Get the balance for a given owner address (which can be a wallet or an object)
19/// and currency type (e.g. 0x2::sui::SUI)
20pub fn get_balance(
21    owner: SuiAddress,
22    child_object_resolver: &dyn ChildObjectResolver,
23    currency_type: TypeTag,
24) -> SuiResult<u64> {
25    let balance_type = Balance::type_tag(currency_type);
26    let address_balance =
27        AccumulatorValue::load(child_object_resolver, None, owner, &balance_type)?
28            .and_then(|b| b.as_u128())
29            .unwrap_or(0);
30
31    let u64_balance = if address_balance > u64::MAX as u128 {
32        // This will not happen with normal currency types which have a max supply of u64::MAX
33        // But you can create "fake" supplies (with no metadata or treasury cap) and overlow
34        // the u64 limit.
35        tracing::warn!(
36            "address balance for {} {} is greater than u64::MAX",
37            owner,
38            balance_type.to_canonical_string(true)
39        );
40        u64::MAX
41    } else {
42        address_balance as u64
43    };
44
45    Ok(u64_balance)
46}
47
48pub trait GetDynamicFieldsIter {
49    fn get_dynamic_fields_iterator(
50        &self,
51        parent: ObjectID,
52        cursor: Option<ObjectID>,
53    ) -> Result<Box<dyn Iterator<Item = DynamicFieldIteratorItem> + '_>, StorageError>;
54}
55
56impl GetDynamicFieldsIter for &IndexStoreTables {
57    fn get_dynamic_fields_iterator(
58        &self,
59        parent: ObjectID,
60        cursor: Option<ObjectID>,
61    ) -> Result<Box<dyn Iterator<Item = DynamicFieldIteratorItem> + '_>, StorageError> {
62        Ok(Box::new(
63            IndexStoreTables::get_dynamic_fields_iterator(self, parent, cursor).map(
64                move |r| match r {
65                    Ok((field_id, _)) => Ok(DynamicFieldKey { parent, field_id }),
66                    Err(e) => Err(e),
67                },
68            ),
69        ))
70    }
71}
72
73impl<T: RpcIndexes> GetDynamicFieldsIter for T {
74    fn get_dynamic_fields_iterator(
75        &self,
76        parent: ObjectID,
77        cursor: Option<ObjectID>,
78    ) -> Result<Box<dyn Iterator<Item = DynamicFieldIteratorItem> + '_>, StorageError> {
79        Ok(Box::new(self.dynamic_field_iter(parent, cursor)?))
80    }
81}
82
83/// Get all currency types for a given owner address (which can be a wallet or an object)
84pub fn get_currency_types_for_owner(
85    owner: SuiAddress,
86    child_object_resolver: &dyn ChildObjectResolver,
87    index_tables: impl GetDynamicFieldsIter,
88    limit: usize,
89    cursor: Option<ObjectID>,
90) -> SuiResult<Vec<TypeTag>> {
91    let Some(owner_obj) = AccumulatorOwner::load_object(child_object_resolver, None, owner)? else {
92        return Ok(Vec::new());
93    };
94
95    let owner_version = owner_obj.version();
96
97    let accumulator_owner_obj = AccumulatorOwner::from_object(owner_obj)?;
98
99    if accumulator_owner_obj.owner != owner {
100        debug_fatal!("owner object owner does not match the requested owner");
101        return Ok(Vec::new());
102    };
103
104    let bag_id = accumulator_owner_obj.balances.id.object_id();
105
106    // get all balance types for the owner
107    let accumulator_metadata: Vec<_> = index_tables
108        .get_dynamic_fields_iterator(*bag_id, cursor)?
109        .take(limit)
110        .collect();
111
112    let mut currency_types = Vec::new();
113    for result in accumulator_metadata {
114        let DynamicFieldKey { parent, field_id } = result?;
115        debug_assert_eq!(parent, *bag_id);
116
117        if let Some(object) =
118            child_object_resolver.read_child_object(bag_id, &field_id, owner_version)?
119        {
120            let ty = object
121                .data
122                .try_as_move()
123                .expect("accumulator metadata object is not a move object")
124                .type_();
125
126            let Some(currency_type) = ty.balance_accumulator_metadata_field_type_maybe() else {
127                // This should currently never happen. But in the future, there may be non-balance
128                // accumulator types, in which case we would need to skip them here.
129                debug_fatal!(
130                    "accumulator metadata object is not a balance accumulator metadata field"
131                );
132                continue;
133            };
134
135            currency_types.push(currency_type);
136        }
137    }
138
139    Ok(currency_types)
140}
141
142/// Get all balances and corresponding currency types for a given owner address
143/// (which can be a wallet or an object)
144pub fn get_all_balances_for_owner(
145    owner: SuiAddress,
146    child_object_resolver: &dyn ChildObjectResolver,
147    index_tables: impl GetDynamicFieldsIter,
148    limit: usize,
149    cursor: Option<ObjectID>,
150) -> SuiResult<Vec<(TypeTag, u64)>> {
151    let currency_types =
152        get_currency_types_for_owner(owner, child_object_resolver, index_tables, limit, cursor)?;
153    let mut balances = Vec::new();
154    for currency_type in currency_types {
155        let balance = get_balance(owner, child_object_resolver, currency_type.clone())?;
156        balances.push((currency_type, balance));
157    }
158    Ok(balances)
159}
sui_core/accumulators/balances.rs

sui_core/accumulators/
balances.rs
