sui_rpc_store/indexer/
balance.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Sequential pipeline that feeds the
5//! [`schema::balance`](crate::schema::balance) CF.
6//!
7//! For every transaction in the checkpoint, call
8//! [`sui_types::balance_change::derive_detailed_balance_changes_2`]
9//! and forward the returned `(coin_amount, address_amount)`
10//! deltas straight into the CF as a single combined merge operand
11//! per `(owner, coin_type)`.
12//!
13//! The `derive_detailed_balance_changes_2` helper already
14//! consolidates input and output coin objects (for the *coin*
15//! side) and parses the effects' accumulator writes (for the
16//! *address* side), so the pipeline doesn't need to walk objects
17//! itself.
18//!
19//! Balances are an address-level aggregate: the helper's coin side
20//! counts only `AddressOwner` and `ConsensusAddressOwner` coins
21//! (combined per address), matching [`Balance::restore`]'s owner
22//! filter below -- the two must agree, or a store restored from a
23//! live snapshot would diverge from one built by tip indexing.
24//! This is the same rule `sui-indexer-alt-consistent-store`'s
25//! `balances` handler documents; object-owned coins stay
26//! discoverable through `object_by_owner` under their parent, they
27//! are just not a balance.
28
29use std::collections::HashMap;
30use std::sync::Arc;
31
32use async_trait::async_trait;
33use move_core_types::language_storage::TypeTag;
34use sui_consistent_store::Batch;
35use sui_consistent_store::Restore;
36use sui_indexer_alt_framework::pipeline::Processor;
37use sui_indexer_alt_framework::pipeline::sequential;
38use sui_types::SUI_ACCUMULATOR_ROOT_OBJECT_ID;
39use sui_types::accumulator_root::AccumulatorKey;
40use sui_types::accumulator_root::AccumulatorValue;
41use sui_types::balance_change::derive_detailed_balance_changes_2;
42use sui_types::base_types::SuiAddress;
43use sui_types::coin::Coin;
44use sui_types::full_checkpoint_content::Checkpoint;
45use sui_types::object::Object;
46use sui_types::object::Owner;
47
48use crate::RpcStoreSchema;
49use crate::indexer::Schema;
50use crate::indexer::Store;
51use crate::schema::balance;
52use crate::schema::balance::Key;
53
54/// Pipeline marker for `balance`.
55pub struct Balance;
56
57#[derive(Debug)]
58pub struct Delta {
59    pub owner: SuiAddress,
60    pub coin_type: TypeTag,
61    /// Change to the coin-derived component (sum of owned
62    /// `Coin<T>` deltas).
63    pub coin: i128,
64    /// Change to the accumulator-derived component (sum of
65    /// per-tx accumulator writes against `(owner, coin_type)`).
66    pub address: i128,
67}
68
69#[async_trait]
70impl Processor for Balance {
71    const NAME: &'static str = "balance";
72    type Value = Delta;
73
74    async fn process(&self, checkpoint: &Arc<Checkpoint>) -> anyhow::Result<Vec<Delta>> {
75        let mut deltas = Vec::new();
76        for tx in &checkpoint.transactions {
77            for change in derive_detailed_balance_changes_2(&tx.effects, &checkpoint.object_set) {
78                deltas.push(Delta {
79                    owner: change.address,
80                    coin_type: change.coin_type,
81                    coin: change.coin_amount,
82                    address: change.address_amount,
83                });
84            }
85        }
86        Ok(deltas)
87    }
88}
89
90impl Restore for Balance {
91    type Schema = RpcStoreSchema;
92
93    /// Stage merge operands derived from a single live object.
94    /// Two sources contribute to a balance row, both recoverable
95    /// from the live object set:
96    ///
97    /// - **Coin half**: address-owned (and consensus-address-owned)
98    ///   `Coin<T>` objects. The coin's `balance` field is credited
99    ///   to the `(owner, coin_type)` row's coin component.
100    ///
101    /// - **Address half**: dynamic-field objects parented to
102    ///   [`SUI_ACCUMULATOR_ROOT_OBJECT_ID`]. These carry the
103    ///   per-`(owner, coin_type)` accumulator balance, which the
104    ///   tip pipeline would otherwise re-derive from
105    ///   `AccumulatorWrite` events.
106    ///
107    /// Everything else (shared / immutable objects, non-coin
108    /// address-owned objects, dynamic fields under other parents)
109    /// contributes no balance row.
110    fn restore(
111        &self,
112        schema: &Self::Schema,
113        object: &Object,
114        batch: &mut Batch,
115    ) -> anyhow::Result<()> {
116        match object.owner() {
117            Owner::AddressOwner(owner) | Owner::ConsensusAddressOwner { owner, .. } => {
118                if let Some((coin_type, value)) = coin_balance_for_restore(object)? {
119                    let (key, val) = balance::delta(*owner, coin_type, value as i128, 0);
120                    batch.merge(&schema.balance, &key, &val)?;
121                }
122            }
123            Owner::ObjectOwner(parent) if *parent == SUI_ACCUMULATOR_ROOT_OBJECT_ID.into() => {
124                if let Some((owner, coin_type, balance_value)) = address_balance_info(object) {
125                    let (key, val) = balance::delta(owner, coin_type, 0, balance_value);
126                    batch.merge(&schema.balance, &key, &val)?;
127                }
128            }
129            _ => {}
130        }
131        Ok(())
132    }
133}
134
135/// Extract the `(coin_type, balance)` pair for a coin object, or
136/// `None` if `object` is not a coin or carries a non-struct type
137/// tag.
138fn coin_balance_for_restore(object: &Object) -> anyhow::Result<Option<(TypeTag, u64)>> {
139    Ok(Coin::extract_balance_if_coin(object)
140        .map_err(|e| anyhow::anyhow!("Failed to deserialize coin object {}: {e}", object.id()))?
141        .and_then(|(type_, value)| match type_ {
142            TypeTag::Struct(struct_tag) => Some((TypeTag::Struct(struct_tag), value)),
143            _ => None,
144        }))
145}
146
147/// Extract `(owner, coin_type, balance)` from a dynamic-field
148/// object parented to the accumulator root. Returns `None` for
149/// non-balance fields, fields whose value cannot be parsed as a
150/// `u128`, or non-positive balances.
151fn address_balance_info(object: &Object) -> Option<(SuiAddress, TypeTag, i128)> {
152    let move_object = object.data.try_as_move()?;
153    let TypeTag::Struct(coin_type) = move_object.type_().balance_accumulator_field_type_maybe()?
154    else {
155        return None;
156    };
157    let (key, value): (AccumulatorKey, AccumulatorValue) = move_object.try_into().ok()?;
158    let balance_value = value.as_u128()? as i128;
159    if balance_value <= 0 {
160        return None;
161    }
162    Some((key.owner, TypeTag::Struct(coin_type), balance_value))
163}
164
165#[async_trait]
166impl sequential::Handler for Balance {
167    type Store = Store;
168    /// Combine deltas observed in this checkpoint by
169    /// `(owner, coin_type)` so a single combined merge operand is
170    /// staged per key instead of many small ones.
171    type Batch = HashMap<Key, (i128, i128)>;
172
173    fn batch(&self, batch: &mut Self::Batch, values: std::vec::IntoIter<Delta>) {
174        for d in values {
175            let entry = batch
176                .entry(Key {
177                    owner: d.owner,
178                    coin_type: d.coin_type,
179                })
180                .or_insert((0, 0));
181            entry.0 = entry.0.saturating_add(d.coin);
182            entry.1 = entry.1.saturating_add(d.address);
183        }
184    }
185
186    async fn commit<'a>(
187        &self,
188        batch: &Self::Batch,
189        conn: &mut sui_consistent_store::Connection<'a, Schema>,
190    ) -> anyhow::Result<usize> {
191        let cf = &conn.store.schema().balance;
192        for (key, (coin, address)) in batch {
193            let (_, value) = balance::delta(key.owner, key.coin_type.clone(), *coin, *address);
194            conn.batch.merge(cf, key, &value)?;
195        }
196        Ok(batch.len())
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use std::sync::Arc;
203
204    use sui_consistent_store::Db;
205    use sui_consistent_store::DbOptions;
206    use sui_types::base_types::ObjectID;
207    use sui_types::test_checkpoint_data_builder::TestCheckpointBuilder;
208
209    use super::*;
210
211    #[tokio::test]
212    async fn process_runs_against_synthetic_checkpoint() {
213        let checkpoint = Arc::new(TestCheckpointBuilder::new(1).build_checkpoint());
214        let _ = Balance.process(&checkpoint).await.unwrap();
215    }
216
217    #[test]
218    fn restore_credits_coin_half_for_address_owned_gas_coin() {
219        let dir = tempfile::tempdir().unwrap();
220        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
221
222        let owner = SuiAddress::ZERO;
223        let coin = Object::with_id_owner_gas_for_testing(ObjectID::from_single_byte(5), owner, 42);
224        let coin_type = coin.coin_type_maybe().unwrap();
225
226        let mut batch = db.batch();
227        Balance.restore(&schema, &coin, &mut batch).unwrap();
228        batch.commit().unwrap();
229
230        let balance = schema
231            .get_balance(owner, coin_type)
232            .unwrap()
233            .expect("balance row present");
234        assert_eq!(balance.coin, 42);
235        // No matching accumulator-root dynamic-field object was
236        // restored alongside the coin, so the address half stays
237        // zero. A test that exercises the address half lives below.
238        assert_eq!(balance.address, 0);
239    }
240
241    /// A coin held by an object (dynamic field, transfer-to-object)
242    /// contributes no balance row -- mirroring the tip pipeline, whose
243    /// `derive_detailed_balance_changes_2` coin side counts only
244    /// address-held coins. If either side ever drifts, restored and
245    /// tip-built stores diverge.
246    #[test]
247    fn restore_skips_object_owned_coins() {
248        use sui_types::object::Owner;
249
250        let dir = tempfile::tempdir().unwrap();
251        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
252
253        let parent = ObjectID::from_single_byte(0x42);
254        let mut coin = Object::with_id_owner_gas_for_testing(
255            ObjectID::from_single_byte(6),
256            SuiAddress::ZERO,
257            42,
258        );
259        coin.owner = Owner::ObjectOwner(parent.into());
260        let coin_type = coin.coin_type_maybe().unwrap();
261
262        let mut batch = db.batch();
263        Balance.restore(&schema, &coin, &mut batch).unwrap();
264        batch.commit().unwrap();
265
266        assert!(
267            schema
268                .get_balance(parent.into(), coin_type)
269                .unwrap()
270                .is_none(),
271            "an object-owned coin must not credit the parent's id as a balance",
272        );
273    }
274
275    #[test]
276    fn restore_skips_non_coin_objects() {
277        let dir = tempfile::tempdir().unwrap();
278        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
279
280        let owner = SuiAddress::ZERO;
281        let non_coin = Object::with_id_owner_for_testing(ObjectID::from_single_byte(9), owner);
282
283        let mut batch = db.batch();
284        Balance.restore(&schema, &non_coin, &mut batch).unwrap();
285        batch.commit().unwrap();
286        // Nothing to assert on read because we don't know the
287        // (non-coin) type to query by; the meaningful assertion
288        // is just that `restore` returned `Ok` without staging a
289        // bad write.
290    }
291}
sui_rpc_store/indexer/balance.rs

sui_rpc_store/indexer/
balance.rs
