1use 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#[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 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
70impl 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#[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 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
210pub 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 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 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 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 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 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}