sui_rpc_store/schema/
object_by_owner.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! `(OwnerKind, type, inverted_balance?, ObjectID)` → latest live
5//! `version`.
6//!
7//! Supports owner-and-type filtering with optional balance-based
8//! ordering. The leading [`OwnerKind`] byte clusters entries by
9//! ownership category: address-owned, object-owned, shared,
10//! immutable. The address variants carry the owning address in the
11//! key; shared and immutable do not. Within each
12//! `(kind, owner, type)` group, coin-like objects carry the
13//! ones-complement of their balance (`!balance`) so richer coins
14//! sort first under a forward prefix scan; non-coin objects carry
15//! no balance at all.
16
17use bytes::Buf;
18use bytes::BufMut;
19use move_core_types::language_storage::StructTag;
20use sui_consistent_store::Decode;
21use sui_consistent_store::Encode;
22use sui_consistent_store::Iter;
23use sui_consistent_store::error::DecodeError;
24use sui_consistent_store::error::EncodeError;
25use sui_consistent_store::error::Error;
26use sui_consistent_store::reader::Reader;
27use sui_types::base_types::ObjectID;
28use sui_types::base_types::SUI_ADDRESS_LENGTH;
29use sui_types::base_types::SuiAddress;
30use sui_types::object::Object;
31use sui_types::object::Owner;
32
33use crate::schema::primitives::U64Varint;
34
35pub const NAME: &str = "object_by_owner";
36
37/// The four kinds of ownership this index distinguishes. The
38/// address-owner and object-owner variants carry the owning
39/// `SuiAddress` inline; shared and immutable owners have no
40/// owning address.
41#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
42pub enum OwnerKind {
43    AddressOwner(SuiAddress),
44    ObjectOwner(SuiAddress),
45    Shared,
46    Immutable,
47}
48
49impl OwnerKind {
50    /// Project a canonical [`Owner`] onto this index's
51    /// [`OwnerKind`] shape.
52    ///
53    /// - `ConsensusAddressOwner` collapses into `AddressOwner` so
54    ///   address-based listings return objects regardless of
55    ///   whether they sit on the consensus path.
56    /// - `Party` is intentionally unimplemented: the canonical
57    ///   type is still in flux upstream.
58    pub fn from_owner(owner: &Owner) -> Self {
59        match owner {
60            Owner::AddressOwner(address) => OwnerKind::AddressOwner(*address),
61            Owner::ObjectOwner(address) => OwnerKind::ObjectOwner(*address),
62            Owner::Shared { .. } => OwnerKind::Shared,
63            Owner::Immutable => OwnerKind::Immutable,
64            Owner::ConsensusAddressOwner { owner, .. } => OwnerKind::AddressOwner(*owner),
65            Owner::Party { .. } => todo!("Party owner WIP"),
66        }
67    }
68}
69
70/// Encoded as `kind_tag(1) || owner?(32)`: the leading bytes shared
71/// by every [`Key`] of this kind. Encoding an `OwnerKind` directly
72/// (or composing it with a [`TypeFilter`](super::type_filter::TypeFilter)
73/// via a tuple) yields a prefix that
74/// [`DbMap::iter_prefix`](sui_consistent_store::DbMap::iter_prefix)
75/// can scan; see [`Key`]'s `Encode` impl, which delegates here for
76/// its leading bytes.
77impl Encode for OwnerKind {
78    fn encode_into<B: BufMut>(&self, buf: &mut B) -> Result<(), EncodeError> {
79        match self {
80            OwnerKind::AddressOwner(addr) => {
81                buf.put_u8(0);
82                buf.put_slice(addr.as_ref());
83            }
84            OwnerKind::ObjectOwner(addr) => {
85                buf.put_u8(1);
86                buf.put_slice(addr.as_ref());
87            }
88            OwnerKind::Shared => buf.put_u8(2),
89            OwnerKind::Immutable => buf.put_u8(3),
90        }
91        Ok(())
92    }
93}
94
95/// Encoded as
96/// `kind_tag(1) || owner?(32) || type(bcs) || balance_tag(1) || balance?(8 BE) || object_id(32)`.
97///
98/// `kind_tag` distinguishes the four owner kinds (`0` =
99/// AddressOwner, `1` = ObjectOwner, `2` = Shared, `3` = Immutable).
100/// The 32-byte owning address follows only for the two owner-kind
101/// variants. `balance_tag` is `0` for non-coin rows (no balance
102/// follows) and `1` for coin-like rows (followed by 8 big-endian
103/// bytes of the ones-complement of the coin's balance —
104/// `!balance`, so richer coins sort first under a forward prefix
105/// scan).
106#[derive(Debug, Clone, PartialEq, Eq, Hash)]
107pub struct Key {
108    pub kind: OwnerKind,
109    pub type_: StructTag,
110    pub inverted_balance: Option<u64>,
111    pub object_id: ObjectID,
112}
113
114pub type Value = U64Varint;
115
116impl Encode for Key {
117    fn encode_into<B: BufMut>(&self, buf: &mut B) -> Result<(), EncodeError> {
118        self.kind.encode_into(buf)?;
119        let type_bytes = bcs::to_bytes(&self.type_)
120            .map_err(|e| EncodeError::with_source("bcs encode StructTag", e))?;
121        buf.put_slice(&type_bytes);
122        match self.inverted_balance {
123            None => buf.put_u8(0),
124            Some(b) => {
125                buf.put_u8(1);
126                buf.put_slice(&b.to_be_bytes());
127            }
128        }
129        buf.put_slice(self.object_id.as_ref());
130        Ok(())
131    }
132}
133
134impl Decode for Key {
135    fn decode<B: Buf>(buf: &mut B) -> Result<Self, DecodeError> {
136        if !buf.has_remaining() {
137            return Err(DecodeError::msg(format!("{NAME} Key empty")));
138        }
139        let kind = match buf.get_u8() {
140            0 => OwnerKind::AddressOwner(read_address(buf)?),
141            1 => OwnerKind::ObjectOwner(read_address(buf)?),
142            2 => OwnerKind::Shared,
143            3 => OwnerKind::Immutable,
144            v => {
145                return Err(DecodeError::msg(format!(
146                    "{NAME} unknown OwnerKind tag: {v}",
147                )));
148            }
149        };
150
151        // Consume one `StructTag`'s worth of bytes via the
152        // streaming BCS parser. The parser stops at the StructTag's
153        // natural end and leaves the rest of the buffer (balance
154        // tag, balance payload, object id) intact.
155        let type_ = crate::schema::primitives::read_struct_tag(buf)?;
156
157        if !buf.has_remaining() {
158            return Err(DecodeError::msg(format!("{NAME} missing balance tag")));
159        }
160        let inverted_balance = match buf.get_u8() {
161            0 => None,
162            1 => {
163                if buf.remaining() < 8 {
164                    return Err(DecodeError::msg(format!("{NAME} missing balance payload",)));
165                }
166                Some(buf.get_u64())
167            }
168            v => {
169                return Err(DecodeError::msg(
170                    format!("{NAME} invalid balance tag: {v}",),
171                ));
172            }
173        };
174
175        if buf.remaining() != ObjectID::LENGTH {
176            return Err(DecodeError::msg(format!(
177                "{NAME} expected {} trailing bytes for object_id, got {}",
178                ObjectID::LENGTH,
179                buf.remaining(),
180            )));
181        }
182        let mut id = [0u8; ObjectID::LENGTH];
183        buf.copy_to_slice(&mut id);
184
185        Ok(Key {
186            kind,
187            type_,
188            inverted_balance,
189            object_id: ObjectID::new(id),
190        })
191    }
192}
193
194fn read_address<B: Buf>(buf: &mut B) -> Result<SuiAddress, DecodeError> {
195    if buf.remaining() < SUI_ADDRESS_LENGTH {
196        return Err(DecodeError::msg(format!(
197            "{NAME} missing owner address: {} bytes left",
198            buf.remaining(),
199        )));
200    }
201    let mut bytes = [0u8; SUI_ADDRESS_LENGTH];
202    buf.copy_to_slice(&mut bytes);
203    SuiAddress::from_bytes(bytes).map_err(|e| DecodeError::with_source("decode SuiAddress", e))
204}
205
206pub fn options(resolver: &sui_consistent_store::CfOptionsResolver) -> rocksdb::Options {
207    resolver.options(NAME)
208}
209
210/// Build the `(Key, Value)` pair indexing a Move object by owner.
211///
212/// Returns `None` for objects that aren't Move objects (packages,
213/// for example) — those have no `StructTag` and aren't part of
214/// this index. For coin-like Move objects the balance is captured
215/// as the ones-complement `!balance` so richer coins sort first
216/// within their `(owner, type)` group.
217pub fn store(object: &Object) -> Option<(Key, U64Varint)> {
218    let type_: StructTag = object.type_()?.clone().into();
219    Some((
220        Key {
221            kind: OwnerKind::from_owner(object.owner()),
222            type_,
223            inverted_balance: object.as_coin_maybe().map(|coin| !coin.balance.value()),
224            object_id: object.id(),
225        },
226        U64Varint(object.version().value()),
227    ))
228}
229
230impl<R: Reader> super::RpcStoreSchema<R> {
231    /// Iterate over every live object owned (in the address-owner
232    /// sense) by `owner`, in the natural sort order of the index:
233    /// by Move type, then within each type by descending balance
234    /// for coin-like objects, then by object id.
235    pub fn iter_objects_owned_by_address(
236        &self,
237        owner: SuiAddress,
238    ) -> Result<Iter<'_, Key, U64Varint>, Error> {
239        self.object_by_owner
240            .iter_prefix(&OwnerKind::AddressOwner(owner))
241    }
242
243    /// Iterate over every live object owned (in the address-owner
244    /// sense) by `owner` whose Move type matches `type_filter`.
245    /// See [`type_filter::TypeFilter`](super::type_filter::TypeFilter)
246    /// for the matching contract.
247    pub fn iter_objects_owned_by_address_of_type(
248        &self,
249        owner: SuiAddress,
250        type_filter: super::type_filter::TypeFilter,
251    ) -> Result<Iter<'_, Key, U64Varint>, Error> {
252        self.object_by_owner
253            .iter_prefix(&(OwnerKind::AddressOwner(owner), type_filter))
254    }
255
256    /// Iterate over every live object owned (in the object-owner
257    /// sense) by the parent object at the given id.
258    pub fn iter_objects_owned_by_object(
259        &self,
260        parent: SuiAddress,
261    ) -> Result<Iter<'_, Key, U64Varint>, Error> {
262        self.object_by_owner
263            .iter_prefix(&OwnerKind::ObjectOwner(parent))
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use std::collections::BTreeSet;
270
271    use sui_consistent_store::Db;
272    use sui_consistent_store::DbOptions;
273
274    use super::*;
275    use crate::RpcStoreSchema;
276
277    fn sui_tag() -> StructTag {
278        StructTag {
279            address: move_core_types::account_address::AccountAddress::new([2u8; 32]),
280            module: move_core_types::identifier::Identifier::new("sui").unwrap(),
281            name: move_core_types::identifier::Identifier::new("SUI").unwrap(),
282            type_params: vec![],
283        }
284    }
285
286    fn round_trip(key: Key) {
287        let bytes = key.encode().unwrap();
288        let decoded = Key::decode(&mut &bytes[..]).unwrap();
289        assert_eq!(key, decoded);
290    }
291
292    #[test]
293    fn address_owner_with_balance_round_trips() {
294        round_trip(Key {
295            kind: OwnerKind::AddressOwner(SuiAddress::from_bytes([1u8; 32]).unwrap()),
296            type_: sui_tag(),
297            inverted_balance: Some(!1_000_000u64),
298            object_id: ObjectID::new([7u8; 32]),
299        });
300    }
301
302    #[test]
303    fn object_owner_without_balance_round_trips() {
304        round_trip(Key {
305            kind: OwnerKind::ObjectOwner(SuiAddress::from_bytes([3u8; 32]).unwrap()),
306            type_: sui_tag(),
307            inverted_balance: None,
308            object_id: ObjectID::new([8u8; 32]),
309        });
310    }
311
312    #[test]
313    fn shared_round_trips() {
314        round_trip(Key {
315            kind: OwnerKind::Shared,
316            type_: sui_tag(),
317            inverted_balance: None,
318            object_id: ObjectID::new([9u8; 32]),
319        });
320    }
321
322    #[test]
323    fn immutable_round_trips() {
324        round_trip(Key {
325            kind: OwnerKind::Immutable,
326            type_: sui_tag(),
327            inverted_balance: None,
328            object_id: ObjectID::new([0xAAu8; 32]),
329        });
330    }
331
332    fn fresh_db() -> (tempfile::TempDir, sui_consistent_store::Db, RpcStoreSchema) {
333        let dir = tempfile::tempdir().unwrap();
334        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
335        (dir, db, schema)
336    }
337
338    fn dummy_object(id: ObjectID, owner: SuiAddress) -> Object {
339        Object::with_id_owner_for_testing(id, owner)
340    }
341
342    #[test]
343    fn store_derives_key_from_object() {
344        let owner = SuiAddress::from_bytes([5u8; 32]).unwrap();
345        let id = ObjectID::random();
346        let object = dummy_object(id, owner);
347        let (key, value) = store(&object).expect("Move object");
348
349        assert_eq!(key.kind, OwnerKind::AddressOwner(owner));
350        assert_eq!(key.object_id, id);
351        assert_eq!(value.0, object.version().value());
352        // Gas-coin test objects carry a balance, so the inverted
353        // balance should be populated.
354        assert!(key.inverted_balance.is_some());
355    }
356
357    #[test]
358    fn iter_returns_empty_for_owner_with_no_objects() {
359        let (_dir, _db, schema) = fresh_db();
360        let owner = SuiAddress::from_bytes([1u8; 32]).unwrap();
361        let count = schema.iter_objects_owned_by_address(owner).unwrap().count();
362        assert_eq!(count, 0);
363    }
364
365    #[test]
366    fn iter_with_type_filter_narrows_to_matching_objects() {
367        // All gas-coin test objects share the same Move type
368        // (`0x2::coin::Coin<0x2::sui::SUI>`). A `TypeFilter::Type`
369        // pointing at that type should return every one of them;
370        // a `TypeFilter::Type` at a different type should return
371        // none.
372        let (_dir, db, schema) = fresh_db();
373        let owner = SuiAddress::from_bytes([1u8; 32]).unwrap();
374
375        let mut expected_ids = BTreeSet::new();
376        let mut batch = db.batch();
377        let mut shared_type = None;
378        for _ in 0..3 {
379            let id = ObjectID::random();
380            expected_ids.insert(id);
381            let (k, v) = store(&dummy_object(id, owner)).unwrap();
382            shared_type.get_or_insert(k.type_.clone());
383            batch.put(&schema.object_by_owner, &k, &v).unwrap();
384        }
385        batch.commit().unwrap();
386
387        let shared_type = shared_type.unwrap();
388        let matching_filter = super::super::type_filter::TypeFilter::Type(shared_type.clone());
389        let found: BTreeSet<ObjectID> = schema
390            .iter_objects_owned_by_address_of_type(owner, matching_filter)
391            .unwrap()
392            .map(|res| res.unwrap().0.object_id)
393            .collect();
394        assert_eq!(found, expected_ids);
395
396        let mismatched_filter = super::super::type_filter::TypeFilter::Type(StructTag {
397            name: move_core_types::identifier::Identifier::new("Other").unwrap(),
398            ..shared_type
399        });
400        let mismatched_count = schema
401            .iter_objects_owned_by_address_of_type(owner, mismatched_filter)
402            .unwrap()
403            .count();
404        assert_eq!(mismatched_count, 0);
405    }
406
407    #[test]
408    fn iter_finds_only_objects_for_target_owner() {
409        let (_dir, db, schema) = fresh_db();
410        let target = SuiAddress::from_bytes([1u8; 32]).unwrap();
411        let other = SuiAddress::from_bytes([2u8; 32]).unwrap();
412
413        let mut target_ids = BTreeSet::new();
414        let mut batch = db.batch();
415        for _ in 0..3 {
416            let id = ObjectID::random();
417            target_ids.insert(id);
418            let (k, v) = store(&dummy_object(id, target)).unwrap();
419            batch.put(&schema.object_by_owner, &k, &v).unwrap();
420        }
421        for _ in 0..2 {
422            let id = ObjectID::random();
423            let (k, v) = store(&dummy_object(id, other)).unwrap();
424            batch.put(&schema.object_by_owner, &k, &v).unwrap();
425        }
426        batch.commit().unwrap();
427
428        let found: BTreeSet<ObjectID> = schema
429            .iter_objects_owned_by_address(target)
430            .unwrap()
431            .map(|res| res.unwrap().0.object_id)
432            .collect();
433        assert_eq!(found, target_ids);
434    }
435}