Skip to main content

sui_move_natives_latest/
test_scenario.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{
5    get_extension, get_extension_mut, get_nth_struct_field, get_tag_and_layouts, legacy_test_cost,
6    object_runtime::{ObjectRuntime, RuntimeResults, object_store::ChildObjectEffects},
7    scratch::ScratchRuntime,
8};
9use better_any::{Tid, TidAble};
10use indexmap::{IndexMap, IndexSet};
11use move_binary_format::errors::{PartialVMError, PartialVMResult};
12use move_binary_format::{safe_assert, safe_unwrap};
13use move_core_types::{
14    account_address::AccountAddress,
15    annotated_value::{MoveFieldLayout, MoveStructLayout, MoveTypeLayout, MoveValue},
16    annotated_visitor as AV,
17    language_storage::StructTag,
18    vm_status::StatusCode,
19};
20use move_vm_runtime::{
21    execution::values::{Vector, VectorSpecialization},
22    natives::{
23        extensions::NativeExtensionMarker,
24        functions::{NativeContext, NativeResult},
25    },
26};
27use move_vm_runtime::{
28    execution::{
29        Type,
30        values::{self, StructRef, Value},
31    },
32    pop_arg,
33};
34use smallvec::smallvec;
35use std::{
36    borrow::Borrow,
37    cell::RefCell,
38    collections::{BTreeMap, BTreeSet, VecDeque},
39};
40use sui_types::{
41    TypeTag,
42    base_types::{MoveObjectType, ObjectID, SequenceNumber, SuiAddress},
43    config,
44    digests::{ObjectDigest, TransactionDigest},
45    dynamic_field::DynamicFieldInfo,
46    execution::DynamicallyLoadedObjectMetadata,
47    id::UID,
48    in_memory_storage::InMemoryStorage,
49    object::{MoveObject, Object, Owner},
50    storage::{BackingPackageStore, PackageObject, RuntimeObjectResolver},
51};
52
53const E_COULD_NOT_GENERATE_EFFECTS: u64 = 0;
54const E_INVALID_SHARED_OR_IMMUTABLE_USAGE: u64 = 1;
55const E_OBJECT_NOT_FOUND_CODE: u64 = 4;
56const E_UNABLE_TO_ALLOCATE_RECEIVING_TICKET: u64 = 5;
57const E_RECEIVING_TICKET_ALREADY_ALLOCATED: u64 = 6;
58const E_UNABLE_TO_DEALLOCATE_RECEIVING_TICKET: u64 = 7;
59
60type Set<K> = IndexSet<K>;
61
62/// An in-memory test store is a thin wrapper around the in-memory storage in a mutex. The mutex
63/// allows this to be used by both the object runtime (for reading) and the test scenario (for
64/// writing) while hiding mutability.
65#[derive(Tid)]
66pub struct InMemoryTestStore(pub RefCell<InMemoryStorage>);
67impl<'a> NativeExtensionMarker<'a> for &'a InMemoryTestStore {}
68
69impl BackingPackageStore for InMemoryTestStore {
70    fn get_package_object(
71        &self,
72        package_id: &ObjectID,
73    ) -> sui_types::error::SuiResult<Option<PackageObject>> {
74        self.0.borrow().get_package_object(package_id)
75    }
76}
77
78impl RuntimeObjectResolver for InMemoryTestStore {
79    fn read_child_object(
80        &self,
81        parent: &ObjectID,
82        child: &ObjectID,
83        child_version_upper_bound: SequenceNumber,
84    ) -> sui_types::error::SuiResult<Option<Object>> {
85        self.0
86            .borrow()
87            .read_child_object(parent, child, child_version_upper_bound)
88    }
89
90    fn get_object_received_at_version(
91        &self,
92        owner: &ObjectID,
93        receiving_object_id: &ObjectID,
94        receive_object_at_version: SequenceNumber,
95        epoch_id: sui_types::committee::EpochId,
96    ) -> sui_types::error::SuiResult<Option<Object>> {
97        self.0.borrow().get_object_received_at_version(
98            owner,
99            receiving_object_id,
100            receive_object_at_version,
101            epoch_id,
102        )
103    }
104}
105
106// This function updates the inventories based on the transfers and deletes that occurred in the
107// transaction
108// native fun end_transaction(): TransactionResult;
109pub fn end_transaction(
110    context: &mut NativeContext,
111    ty_args: Vec<Type>,
112    args: VecDeque<Value>,
113) -> PartialVMResult<NativeResult> {
114    safe_assert!(ty_args.is_empty());
115    safe_assert!(args.is_empty());
116    // scratch is per-transaction. A single set of native extensions is reused across the
117    // transactions simulated by `test_scenario`, so clear the store here to mirror the fresh
118    // `ScratchRuntime` that real execution installs per transaction.
119    get_extension_mut!(context, ScratchRuntime)?.clear();
120    let object_runtime_ref: &mut ObjectRuntime = get_extension_mut!(context)?;
121    let taken_shared_or_imm: BTreeMap<_, _> = object_runtime_ref
122        .test_inventories
123        .taken
124        .iter()
125        .filter(|(_id, owner)| matches!(owner, Owner::Shared { .. } | Owner::Immutable))
126        .map(|(id, owner)| (*id, owner.clone()))
127        .collect();
128    // set to true if a shared or imm object was:
129    // - transferred in a way that changes it from its original shared/imm state
130    // - wraps the object
131    // if true, we will "abort"
132    let mut incorrect_shared_or_imm_handling = false;
133
134    // Handle the allocated tickets:
135    // * Remove all allocated_tickets in the test inventories.
136    // * For each allocated ticket, if the ticket's object ID is loaded, move it to `received`.
137    // * Otherwise re-insert the allocated ticket into the objects inventory, and mark it to be
138    //   removed from the backing storage (deferred due to needing to have access to `context` which
139    //   has outstanding references at this point).
140    let allocated_tickets =
141        std::mem::take(&mut object_runtime_ref.test_inventories.allocated_tickets);
142    let mut received = BTreeMap::new();
143    let mut unreceived = BTreeSet::new();
144    let loaded_runtime_objects = object_runtime_ref.loaded_runtime_objects();
145    for (id, (metadata, value)) in allocated_tickets {
146        if loaded_runtime_objects.contains_key(&id) {
147            received.insert(id, metadata);
148        } else {
149            unreceived.insert(id);
150            // This must be untouched since the allocated ticket is still live, so ok to re-insert.
151            object_runtime_ref
152                .test_inventories
153                .objects
154                .insert(id, value);
155        }
156    }
157
158    let object_runtime_state = object_runtime_ref.take_state();
159    // Determine writes and deletes
160    // We pass the received objects since they should be viewed as "loaded" for the purposes of
161    // calculating the effects of the transaction.
162    let results = object_runtime_state.finish(received, ChildObjectEffects::new());
163    let RuntimeResults {
164        writes,
165        user_events,
166        loaded_child_objects: _,
167        created_object_ids,
168        deleted_object_ids,
169        accumulator_events: _,
170        settlement_input_sui: _,
171        settlement_output_sui: _,
172    } = match results {
173        Ok(res) => res,
174        Err(_) => {
175            return Ok(NativeResult::err(
176                legacy_test_cost(),
177                E_COULD_NOT_GENERATE_EFFECTS,
178            ));
179        }
180    };
181    let object_runtime_ref: &mut ObjectRuntime = get_extension_mut!(context)?;
182    let all_active_child_objects_with_values = object_runtime_ref
183        .all_active_child_objects()
184        .filter(|child| child.copied_value.is_some())
185        .map(|child| *child.id)
186        .collect::<BTreeSet<_>>();
187    let inventories = &mut object_runtime_ref.test_inventories;
188    let mut new_object_values = IndexMap::new();
189    let mut transferred = vec![];
190    // cleanup inventories
191    // we will remove all changed objects
192    // - deleted objects need to be removed to mark deletions
193    // - written objects are removed and later replaced to mark new values and new owners
194    // - child objects will not be reflected in transfers, but need to be no longer retrievable
195    for id in deleted_object_ids
196        .iter()
197        .chain(writes.keys())
198        .chain(&all_active_child_objects_with_values)
199    {
200        for addr_inventory in inventories.address_inventories.values_mut() {
201            for s in addr_inventory.values_mut() {
202                s.shift_remove(id);
203            }
204        }
205        for s in &mut inventories.shared_inventory.values_mut() {
206            s.shift_remove(id);
207        }
208        for s in &mut inventories.immutable_inventory.values_mut() {
209            s.shift_remove(id);
210        }
211        inventories.taken.remove(id);
212    }
213
214    // handle transfers, inserting transferred/written objects into their respective inventory
215    let mut created = vec![];
216    let mut written = vec![];
217    for (id, (owner, ty, value)) in writes {
218        // write configs to cache
219        new_object_values.insert(id, (ty.clone(), value.copy_value()));
220        transferred.push((id, owner.clone()));
221        incorrect_shared_or_imm_handling = incorrect_shared_or_imm_handling
222            || taken_shared_or_imm
223                .get(&id)
224                .map(|shared_or_imm_owner| shared_or_imm_owner != &owner)
225                .unwrap_or(/* not incorrect */ false);
226        if created_object_ids.contains(&id) {
227            created.push(id);
228        } else {
229            written.push(id);
230        }
231        match owner {
232            Owner::AddressOwner(a) => {
233                inventories
234                    .address_inventories
235                    .entry(a)
236                    .or_default()
237                    .entry(ty)
238                    .or_default()
239                    .insert(id);
240            }
241            Owner::ObjectOwner(_) => (),
242            Owner::Shared { .. } => {
243                inventories
244                    .shared_inventory
245                    .entry(ty)
246                    .or_default()
247                    .insert(id);
248            }
249            Owner::Immutable => {
250                inventories
251                    .immutable_inventory
252                    .entry(ty)
253                    .or_default()
254                    .insert(id);
255            }
256            Owner::ConsensusAddressOwner { owner, .. } => {
257                inventories
258                    .address_inventories
259                    .entry(owner)
260                    .or_default()
261                    .entry(ty)
262                    .or_default()
263                    .insert(id);
264            }
265            Owner::Party { .. } => {
266                // TODO(Party WIP)
267                todo!("Party WIP")
268            }
269        }
270    }
271
272    // For any unused allocated tickets, remove them from the store.
273    let store: &&InMemoryTestStore = get_extension!(context)?;
274    for id in unreceived {
275        if store.0.borrow_mut().remove_object(id).is_none() {
276            return Ok(NativeResult::err(
277                context.gas_used(),
278                E_UNABLE_TO_DEALLOCATE_RECEIVING_TICKET,
279            ));
280        }
281    }
282
283    // deletions already handled above, but we drop the delete kind for the effects
284    let mut deleted = vec![];
285    for id in deleted_object_ids {
286        // Mark as "incorrect" if a imm object was deleted. Allow shared objects to be deleted though.
287        incorrect_shared_or_imm_handling = incorrect_shared_or_imm_handling
288            || taken_shared_or_imm
289                .get(&id)
290                .is_some_and(|owner| matches!(owner, Owner::Immutable));
291        deleted.push(id);
292    }
293    // find all wrapped objects
294    let mut all_wrapped = BTreeSet::new();
295    let object_runtime_ref: &ObjectRuntime = get_extension!(context)?;
296    find_all_wrapped_objects(
297        context,
298        &mut all_wrapped,
299        new_object_values
300            .iter()
301            .map(|(id, (ty, value))| (id, ty, value)),
302    )?;
303    find_all_wrapped_objects(
304        context,
305        &mut all_wrapped,
306        object_runtime_ref
307            .all_active_child_objects()
308            .filter_map(|child| Some((child.id, child.ty, child.copied_value?))),
309    )?;
310    // mark as "incorrect" if a shared/imm object was wrapped or is a child object
311    incorrect_shared_or_imm_handling = incorrect_shared_or_imm_handling
312        || taken_shared_or_imm.keys().any(|id| {
313            all_wrapped.contains(id) || all_active_child_objects_with_values.contains(id)
314        });
315    // if incorrect handling, return with an 'abort'
316    if incorrect_shared_or_imm_handling {
317        return Ok(NativeResult::err(
318            legacy_test_cost(),
319            E_INVALID_SHARED_OR_IMMUTABLE_USAGE,
320        ));
321    }
322
323    // mark all wrapped as deleted
324    for wrapped in all_wrapped {
325        deleted.push(wrapped)
326    }
327
328    // new input objects are remaining taken objects not written/deleted
329    let object_runtime_ref: &mut ObjectRuntime = get_extension_mut!(context)?;
330    let mut config_settings = vec![];
331    for child in object_runtime_ref.all_active_child_objects() {
332        let s: StructTag = child.ty.clone().into();
333        let is_setting = DynamicFieldInfo::is_dynamic_field(&s)
334            && matches!(&s.type_params[1], TypeTag::Struct(s) if config::is_setting(s));
335        if is_setting {
336            config_settings.push((
337                *child.owner,
338                *child.id,
339                child.ty.clone(),
340                child.copied_value,
341            ));
342        }
343    }
344    for (config, setting, ty, value) in config_settings {
345        object_runtime_ref.config_setting_cache_update(config, setting, ty, value)
346    }
347    object_runtime_ref.state.input_objects = object_runtime_ref
348        .test_inventories
349        .taken
350        .iter()
351        .map(|(id, owner)| (*id, owner.clone()))
352        .collect::<BTreeMap<_, _>>();
353    // update inventories
354    // check for bad updates to immutable values
355    for (id, (ty, value)) in new_object_values {
356        debug_assert!(!all_active_child_objects_with_values.contains(&id));
357        if let Some(prev_value) = object_runtime_ref
358            .test_inventories
359            .taken_immutable_values
360            .get(&ty)
361            .and_then(|values| values.get(&id))
362            && !value.equals(prev_value)?
363        {
364            return Ok(NativeResult::err(
365                legacy_test_cost(),
366                E_INVALID_SHARED_OR_IMMUTABLE_USAGE,
367            ));
368        }
369        object_runtime_ref
370            .test_inventories
371            .objects
372            .insert(id, value);
373    }
374    // remove deleted
375    for id in &deleted {
376        object_runtime_ref.test_inventories.objects.remove(id);
377    }
378    // remove active child objects
379    for id in all_active_child_objects_with_values {
380        object_runtime_ref.test_inventories.objects.remove(&id);
381    }
382
383    let effects = transaction_effects(
384        created,
385        written,
386        deleted,
387        transferred,
388        user_events.len() as u64,
389        // TODO: do we need accumulator events here?
390    )?;
391    Ok(NativeResult::ok(legacy_test_cost(), smallvec![effects]))
392}
393
394// native fun take_from_address_by_id<T: key>(account: address, id: ID): T;
395pub fn take_from_address_by_id(
396    context: &mut NativeContext,
397    ty_args: Vec<Type>,
398    mut args: VecDeque<Value>,
399) -> PartialVMResult<NativeResult> {
400    let specified_ty = get_specified_ty(ty_args)?;
401    let id = pop_id(&mut args)?;
402    let account: SuiAddress = pop_arg!(args, AccountAddress).into();
403    pop_arg!(args, StructRef);
404    safe_assert!(args.is_empty());
405    let specified_obj_ty = object_type_of_type(context, &specified_ty)?;
406    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
407    let inventories = &mut object_runtime.test_inventories;
408    let res = take_from_inventory(
409        |x| {
410            inventories
411                .address_inventories
412                .get(&account)
413                .and_then(|inv| inv.get(&specified_obj_ty))
414                .map(|s| s.contains(x))
415                .unwrap_or(false)
416        },
417        &inventories.objects,
418        &mut inventories.taken,
419        &mut object_runtime.state.input_objects,
420        id,
421        Owner::AddressOwner(account),
422    );
423    Ok(match res {
424        Ok(value) => NativeResult::ok(legacy_test_cost(), smallvec![value]),
425        Err(native_err) => native_err,
426    })
427}
428
429// native fun ids_for_address<T: key>(account: address): vector<ID>;
430pub fn ids_for_address(
431    context: &mut NativeContext,
432    ty_args: Vec<Type>,
433    mut args: VecDeque<Value>,
434) -> PartialVMResult<NativeResult> {
435    let specified_ty = get_specified_ty(ty_args)?;
436    let account: SuiAddress = pop_arg!(args, AccountAddress).into();
437    safe_assert!(args.is_empty());
438    let specified_obj_ty = object_type_of_type(context, &specified_ty)?;
439    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
440    let inventories = &mut object_runtime.test_inventories;
441    let ids = inventories
442        .address_inventories
443        .get(&account)
444        .and_then(|inv| inv.get(&specified_obj_ty))
445        .map(|s| s.iter().map(|id| pack_id(*id)).collect::<Vec<Value>>())
446        .unwrap_or_default();
447    let ids_vector = safe_unwrap!(Vector::pack(VectorSpecialization::Container, ids));
448    Ok(NativeResult::ok(legacy_test_cost(), smallvec![ids_vector]))
449}
450
451// native fun most_recent_id_for_address<T: key>(account: address): Option<ID>;
452pub fn most_recent_id_for_address(
453    context: &mut NativeContext,
454    ty_args: Vec<Type>,
455    mut args: VecDeque<Value>,
456) -> PartialVMResult<NativeResult> {
457    let specified_ty = get_specified_ty(ty_args)?;
458    let account: SuiAddress = pop_arg!(args, AccountAddress).into();
459    safe_assert!(args.is_empty());
460    let specified_obj_ty = object_type_of_type(context, &specified_ty)?;
461    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
462    let inventories = &mut object_runtime.test_inventories;
463    let most_recent_id = match inventories.address_inventories.get(&account) {
464        None => pack_option(vector_specialization(&specified_ty), None)?,
465        Some(inv) => most_recent_at_ty(&inventories.taken, inv, &specified_ty, specified_obj_ty)?,
466    };
467    Ok(NativeResult::ok(
468        legacy_test_cost(),
469        smallvec![most_recent_id],
470    ))
471}
472
473// native fun was_taken_from_address(account: address, id: ID): bool;
474pub fn was_taken_from_address(
475    context: &mut NativeContext,
476    ty_args: Vec<Type>,
477    mut args: VecDeque<Value>,
478) -> PartialVMResult<NativeResult> {
479    safe_assert!(ty_args.is_empty());
480    let id = pop_id(&mut args)?;
481    let account: SuiAddress = pop_arg!(args, AccountAddress).into();
482    safe_assert!(args.is_empty());
483    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
484    let inventories = &mut object_runtime.test_inventories;
485    let was_taken = inventories
486        .taken
487        .get(&id)
488        .map(|owner| owner == &Owner::AddressOwner(account))
489        .unwrap_or(false);
490    Ok(NativeResult::ok(
491        legacy_test_cost(),
492        smallvec![Value::bool(was_taken)],
493    ))
494}
495
496// native fun take_immutable_by_id<T: key>(id: ID): T;
497pub fn take_immutable_by_id(
498    context: &mut NativeContext,
499    ty_args: Vec<Type>,
500    mut args: VecDeque<Value>,
501) -> PartialVMResult<NativeResult> {
502    let specified_ty = get_specified_ty(ty_args)?;
503    let id = pop_id(&mut args)?;
504    pop_arg!(args, StructRef);
505    safe_assert!(args.is_empty());
506    let specified_obj_ty = object_type_of_type(context, &specified_ty)?;
507    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
508    let inventories = &mut object_runtime.test_inventories;
509    let res = take_from_inventory(
510        |x| {
511            inventories
512                .immutable_inventory
513                .get(&specified_obj_ty)
514                .map(|s| s.contains(x))
515                .unwrap_or(false)
516        },
517        &inventories.objects,
518        &mut inventories.taken,
519        &mut object_runtime.state.input_objects,
520        id,
521        Owner::Immutable,
522    );
523    Ok(match res {
524        Ok(value) => {
525            inventories
526                .taken_immutable_values
527                .entry(specified_obj_ty)
528                .or_default()
529                .insert(id, value.copy_value());
530            NativeResult::ok(legacy_test_cost(), smallvec![value])
531        }
532        Err(native_err) => native_err,
533    })
534}
535
536// native fun most_recent_immutable_id<T: key>(): Option<ID>;
537pub fn most_recent_immutable_id(
538    context: &mut NativeContext,
539    ty_args: Vec<Type>,
540    args: VecDeque<Value>,
541) -> PartialVMResult<NativeResult> {
542    let specified_ty = get_specified_ty(ty_args)?;
543    safe_assert!(args.is_empty());
544    let specified_obj_ty = object_type_of_type(context, &specified_ty)?;
545    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
546    let inventories = &mut object_runtime.test_inventories;
547    let most_recent_id = most_recent_at_ty(
548        &inventories.taken,
549        &inventories.immutable_inventory,
550        &specified_ty,
551        specified_obj_ty,
552    )?;
553    Ok(NativeResult::ok(
554        legacy_test_cost(),
555        smallvec![most_recent_id],
556    ))
557}
558
559// native fun was_taken_immutable(id: ID): bool;
560pub fn was_taken_immutable(
561    context: &mut NativeContext,
562    ty_args: Vec<Type>,
563    mut args: VecDeque<Value>,
564) -> PartialVMResult<NativeResult> {
565    safe_assert!(ty_args.is_empty());
566    let id = pop_id(&mut args)?;
567    safe_assert!(args.is_empty());
568    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
569    let inventories = &mut object_runtime.test_inventories;
570    let was_taken = inventories
571        .taken
572        .get(&id)
573        .map(|owner| owner == &Owner::Immutable)
574        .unwrap_or(false);
575    Ok(NativeResult::ok(
576        legacy_test_cost(),
577        smallvec![Value::bool(was_taken)],
578    ))
579}
580
581// native fun take_shared_by_id<T: key>(id: ID): T;
582pub fn take_shared_by_id(
583    context: &mut NativeContext,
584    ty_args: Vec<Type>,
585    mut args: VecDeque<Value>,
586) -> PartialVMResult<NativeResult> {
587    let specified_ty = get_specified_ty(ty_args)?;
588    let id = pop_id(&mut args)?;
589    pop_arg!(args, StructRef);
590    safe_assert!(args.is_empty());
591    let specified_obj_ty = object_type_of_type(context, &specified_ty)?;
592    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
593    let inventories = &mut object_runtime.test_inventories;
594    let res = take_from_inventory(
595        |x| {
596            inventories
597                .shared_inventory
598                .get(&specified_obj_ty)
599                .map(|s| s.contains(x))
600                .unwrap_or(false)
601        },
602        &inventories.objects,
603        &mut inventories.taken,
604        &mut object_runtime.state.input_objects,
605        id,
606        Owner::Shared { initial_shared_version: /* dummy */ SequenceNumber::new() },
607    );
608    Ok(match res {
609        Ok(value) => NativeResult::ok(legacy_test_cost(), smallvec![value]),
610        Err(native_err) => native_err,
611    })
612}
613
614// native fun most_recent_id_shared<T: key>(): Option<ID>;
615pub fn most_recent_id_shared(
616    context: &mut NativeContext,
617    ty_args: Vec<Type>,
618    args: VecDeque<Value>,
619) -> PartialVMResult<NativeResult> {
620    let specified_ty = get_specified_ty(ty_args)?;
621    safe_assert!(args.is_empty());
622    let specified_obj_ty = object_type_of_type(context, &specified_ty)?;
623    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
624    let inventories = &mut object_runtime.test_inventories;
625    let most_recent_id = most_recent_at_ty(
626        &inventories.taken,
627        &inventories.shared_inventory,
628        &specified_ty,
629        specified_obj_ty,
630    )?;
631    Ok(NativeResult::ok(
632        legacy_test_cost(),
633        smallvec![most_recent_id],
634    ))
635}
636
637// native fun was_taken_shared(id: ID): bool;
638pub fn was_taken_shared(
639    context: &mut NativeContext,
640    ty_args: Vec<Type>,
641    mut args: VecDeque<Value>,
642) -> PartialVMResult<NativeResult> {
643    safe_assert!(ty_args.is_empty());
644    let id = pop_id(&mut args)?;
645    safe_assert!(args.is_empty());
646    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
647    let inventories = &mut object_runtime.test_inventories;
648    let was_taken = inventories
649        .taken
650        .get(&id)
651        .map(|owner| matches!(owner, Owner::Shared { .. }))
652        .unwrap_or(false);
653    Ok(NativeResult::ok(
654        legacy_test_cost(),
655        smallvec![Value::bool(was_taken)],
656    ))
657}
658
659pub fn allocate_receiving_ticket_for_object(
660    context: &mut NativeContext,
661    ty_args: Vec<Type>,
662    mut args: VecDeque<Value>,
663) -> PartialVMResult<NativeResult> {
664    let ty = get_specified_ty(ty_args)?;
665    let id = pop_id(&mut args)?;
666
667    let abilities = context.type_to_abilities(&ty)?;
668    let Some((tag, layout, _)) = get_tag_and_layouts(context, &ty)? else {
669        return Ok(NativeResult::err(
670            context.gas_used(),
671            E_UNABLE_TO_ALLOCATE_RECEIVING_TICKET,
672        ));
673    };
674    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
675    let object_version = SequenceNumber::new();
676    let inventories = &mut object_runtime.test_inventories;
677    if inventories.allocated_tickets.contains_key(&id) {
678        return Ok(NativeResult::err(
679            context.gas_used(),
680            E_RECEIVING_TICKET_ALREADY_ALLOCATED,
681        ));
682    }
683
684    let obj_value = safe_unwrap!(inventories.objects.remove(&id));
685    let Some(bytes) = obj_value.typed_serialize(&layout) else {
686        return Ok(NativeResult::err(
687            context.gas_used(),
688            E_UNABLE_TO_ALLOCATE_RECEIVING_TICKET,
689        ));
690    };
691    let has_public_transfer = abilities.has_store();
692    let move_object = safe_unwrap!(unsafe {
693        MoveObject::new_from_execution_with_limit(
694            tag.into(),
695            has_public_transfer,
696            object_version,
697            bytes,
698            250 * 1024,
699        )
700    });
701
702    let Some((owner, _)) = inventories
703        .address_inventories
704        .iter()
705        .find(|(_addr, objs)| objs.iter().any(|(_, ids)| ids.contains(&id)))
706    else {
707        return Ok(NativeResult::err(
708            context.gas_used(),
709            E_OBJECT_NOT_FOUND_CODE,
710        ));
711    };
712
713    inventories.allocated_tickets.insert(
714        id,
715        (
716            DynamicallyLoadedObjectMetadata {
717                version: SequenceNumber::new(),
718                digest: ObjectDigest::MIN,
719                owner: Owner::AddressOwner(*owner),
720                storage_rebate: 0,
721                previous_transaction: TransactionDigest::default(),
722            },
723            obj_value,
724        ),
725    );
726
727    let object = Object::new_move(
728        move_object,
729        Owner::AddressOwner(*owner),
730        TransactionDigest::default(),
731    );
732
733    // NB: Must be a `&&` reference since the extension stores a static ref to the object storage.
734    let store: &&InMemoryTestStore = get_extension!(context)?;
735    store.0.borrow_mut().insert_object(object);
736
737    Ok(NativeResult::ok(
738        legacy_test_cost(),
739        smallvec![Value::u64(object_version.value())],
740    ))
741}
742
743pub fn deallocate_receiving_ticket_for_object(
744    context: &mut NativeContext,
745    _ty_args: Vec<Type>,
746    mut args: VecDeque<Value>,
747) -> PartialVMResult<NativeResult> {
748    let id = pop_id(&mut args)?;
749
750    let object_runtime: &mut ObjectRuntime = get_extension_mut!(context)?;
751    let inventories = &mut object_runtime.test_inventories;
752    // Deallocate the ticket -- we should never hit this scenario
753    let Some((_, value)) = inventories.allocated_tickets.remove(&id) else {
754        return Ok(NativeResult::err(
755            context.gas_used(),
756            E_UNABLE_TO_DEALLOCATE_RECEIVING_TICKET,
757        ));
758    };
759
760    // Insert the object value that we saved from earlier and put it back into the object set.
761    // This is fine since it can't have been touched.
762    inventories.objects.insert(id, value);
763
764    // Remove the object from storage. We should never hit this scenario either.
765    let store: &&InMemoryTestStore = get_extension!(context)?;
766    if store.0.borrow_mut().remove_object(id).is_none() {
767        return Ok(NativeResult::err(
768            context.gas_used(),
769            E_UNABLE_TO_DEALLOCATE_RECEIVING_TICKET,
770        ));
771    };
772
773    Ok(NativeResult::ok(legacy_test_cost(), smallvec![]))
774}
775
776// impls
777
778fn take_from_inventory(
779    is_in_inventory: impl FnOnce(&ObjectID) -> bool,
780    objects: &BTreeMap<ObjectID, Value>,
781    taken: &mut BTreeMap<ObjectID, Owner>,
782    input_objects: &mut BTreeMap<ObjectID, Owner>,
783    id: ObjectID,
784    owner: Owner,
785) -> Result<Value, NativeResult> {
786    let is_taken = taken.contains_key(&id);
787    let obj = match objects.get(&id) {
788        Some(obj) if !is_taken && is_in_inventory(&id) => obj,
789        _ => {
790            return Err(NativeResult::err(
791                legacy_test_cost(),
792                E_OBJECT_NOT_FOUND_CODE,
793            ));
794        }
795    };
796    taken.insert(id, owner.clone());
797    input_objects.insert(id, owner);
798    Ok(obj.copy_value())
799}
800
801fn vector_specialization(ty: &Type) -> VectorSpecialization {
802    match ty.try_into() {
803        Ok(s) => s,
804        Err(_) => {
805            debug_assert!(false, "Invalid vector specialization");
806            VectorSpecialization::Container
807        }
808    }
809}
810
811fn most_recent_at_ty(
812    taken: &BTreeMap<ObjectID, Owner>,
813    inv: &BTreeMap<MoveObjectType, Set<ObjectID>>,
814    runtime_ty: &Type,
815    ty: MoveObjectType,
816) -> PartialVMResult<Value> {
817    pack_option(
818        vector_specialization(runtime_ty),
819        most_recent_at_ty_opt(taken, inv, ty),
820    )
821}
822
823fn most_recent_at_ty_opt(
824    taken: &BTreeMap<ObjectID, Owner>,
825    inv: &BTreeMap<MoveObjectType, Set<ObjectID>>,
826    ty: MoveObjectType,
827) -> Option<Value> {
828    let s = inv.get(&ty)?;
829    let most_recent_id = s.iter().rfind(|id| !taken.contains_key(id))?;
830    Some(pack_id(*most_recent_id))
831}
832
833fn get_specified_ty(mut ty_args: Vec<Type>) -> PartialVMResult<Type> {
834    safe_assert!(ty_args.len() == 1);
835    Ok(safe_unwrap!(ty_args.pop()))
836}
837
838// helpers
839fn pop_id(args: &mut VecDeque<Value>) -> PartialVMResult<ObjectID> {
840    let v = match args.pop_back() {
841        None => {
842            return Err(PartialVMError::new(
843                StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR,
844            ));
845        }
846        Some(v) => v,
847    };
848    Ok(get_nth_struct_field(v, 0)?
849        .value_as::<AccountAddress>()?
850        .into())
851}
852
853fn pack_id(a: impl Into<AccountAddress>) -> Value {
854    Value::struct_(values::Struct::pack(vec![Value::address(a.into())]))
855}
856
857fn pack_ids(items: impl IntoIterator<Item = impl Into<AccountAddress>>) -> PartialVMResult<Value> {
858    Vector::pack(
859        VectorSpecialization::Container,
860        items.into_iter().map(pack_id),
861    )
862}
863
864fn pack_vec_map(items: impl IntoIterator<Item = (Value, Value)>) -> PartialVMResult<Value> {
865    Ok(Value::struct_(values::Struct::pack(vec![Vector::pack(
866        VectorSpecialization::Container,
867        items
868            .into_iter()
869            .map(|(k, v)| Value::struct_(values::Struct::pack(vec![k, v]))),
870    )?])))
871}
872
873fn transaction_effects(
874    created: impl IntoIterator<Item = impl Into<AccountAddress>>,
875    written: impl IntoIterator<Item = impl Into<AccountAddress>>,
876    deleted: impl IntoIterator<Item = impl Into<AccountAddress>>,
877    transferred: impl IntoIterator<Item = (ObjectID, Owner)>,
878    num_events: u64,
879) -> PartialVMResult<Value> {
880    let mut transferred_to_account = vec![];
881    let mut transferred_to_object = vec![];
882    let mut shared = vec![];
883    let mut frozen = vec![];
884    for (id, owner) in transferred {
885        match owner {
886            Owner::AddressOwner(a) => {
887                transferred_to_account.push((pack_id(id), Value::address(a.into())))
888            }
889            Owner::ObjectOwner(o) => transferred_to_object.push((pack_id(id), pack_id(o))),
890            Owner::Shared { .. } => shared.push(id),
891            Owner::Immutable => frozen.push(id),
892            Owner::ConsensusAddressOwner { owner, .. } => {
893                transferred_to_account.push((pack_id(id), Value::address(owner.into())))
894            }
895            Owner::Party { .. } => {
896                // TODO(Party WIP)
897                todo!("Party WIP")
898            }
899        }
900    }
901
902    let created_field = pack_ids(created)?;
903    let written_field = pack_ids(written)?;
904    let deleted_field = pack_ids(deleted)?;
905    let transferred_to_account_field = pack_vec_map(transferred_to_account)?;
906    let transferred_to_object_field = pack_vec_map(transferred_to_object)?;
907    let shared_field = pack_ids(shared)?;
908    let frozen_field = pack_ids(frozen)?;
909    let num_events_field = Value::u64(num_events);
910    Ok(Value::struct_(values::Struct::pack(vec![
911        created_field,
912        written_field,
913        deleted_field,
914        transferred_to_account_field,
915        transferred_to_object_field,
916        shared_field,
917        frozen_field,
918        num_events_field,
919    ])))
920}
921
922fn object_type_of_type(context: &NativeContext, ty: &Type) -> PartialVMResult<MoveObjectType> {
923    let TypeTag::Struct(s_tag) = context.type_to_type_tag(ty)? else {
924        return Err(PartialVMError::new(
925            StatusCode::UNKNOWN_INVARIANT_VIOLATION_ERROR,
926        ));
927    };
928    Ok(MoveObjectType::from(*s_tag))
929}
930
931fn pack_option(specialization: VectorSpecialization, opt: Option<Value>) -> PartialVMResult<Value> {
932    let item = match opt {
933        Some(v) => vec![v],
934        None => vec![],
935    };
936    Ok(Value::struct_(values::Struct::pack(vec![Vector::pack(
937        specialization,
938        item,
939    )?])))
940}
941
942fn find_all_wrapped_objects<'a, 'i>(
943    context: &NativeContext,
944    ids: &'i mut BTreeSet<ObjectID>,
945    new_object_values: impl IntoIterator<Item = (&'a ObjectID, &'a MoveObjectType, impl Borrow<Value>)>,
946) -> PartialVMResult<()> {
947    #[derive(Copy, Clone)]
948    enum LookingFor {
949        Wrapped,
950        Uid,
951        Address,
952    }
953
954    struct Traversal<'i, 'u> {
955        state: LookingFor,
956        ids: &'i mut BTreeSet<ObjectID>,
957        uid: &'u MoveStructLayout,
958    }
959
960    impl<'b, 'l> AV::Traversal<'b, 'l> for Traversal<'_, '_> {
961        type Error = AV::Error;
962
963        fn traverse_struct(
964            &mut self,
965            driver: &mut AV::StructDriver<'_, 'b, 'l>,
966        ) -> Result<(), Self::Error> {
967            match self.state {
968                // We're at the top-level of the traversal, looking for an object to recurse into.
969                // We can unconditionally switch to looking for UID fields at the level below,
970                // because we know that all the top-level values are objects.
971                LookingFor::Wrapped => {
972                    while driver
973                        .next_field(&mut Traversal {
974                            state: LookingFor::Uid,
975                            ids: self.ids,
976                            uid: self.uid,
977                        })?
978                        .is_some()
979                    {}
980                }
981
982                // We are looking for UID fields. If we find one (which we confirm by checking its
983                // layout), switch to looking for addresses in its sub-structure.
984                LookingFor::Uid => {
985                    while let Some(MoveFieldLayout { name: _, layout }) = driver.peek_field() {
986                        if matches!(layout, MoveTypeLayout::Struct(s) if s.as_ref() == self.uid) {
987                            driver.next_field(&mut Traversal {
988                                state: LookingFor::Address,
989                                ids: self.ids,
990                                uid: self.uid,
991                            })?;
992                        } else {
993                            driver.next_field(self)?;
994                        }
995                    }
996                }
997
998                // When looking for addresses, recurse through structs, as the address is nested
999                // within the UID.
1000                LookingFor::Address => while driver.next_field(self)?.is_some() {},
1001            }
1002
1003            Ok(())
1004        }
1005
1006        fn traverse_address(
1007            &mut self,
1008            _: &AV::ValueDriver<'_, 'b, 'l>,
1009            address: AccountAddress,
1010        ) -> Result<(), Self::Error> {
1011            // If we're looking for addresses, and we found one, then save it.
1012            if matches!(self.state, LookingFor::Address) {
1013                self.ids.insert(address.into());
1014            }
1015            Ok(())
1016        }
1017    }
1018
1019    let uid = UID::layout();
1020    for (_id, ty, value) in new_object_values {
1021        let type_tag = TypeTag::from(ty.clone());
1022        // NB: We can get the layout from the VM's cache since the types and modules
1023        // associated with all of these types must be in the type/module cache in the VM -- THIS IS
1024        // BECAUSE WE ARE IN TEST SCENARIO ONLY AND THIS MAY NOT GENERALLY HOLD IN A
1025        // MULTI-TRANSACTION SETTING.
1026        let Some(layout) = context.type_tag_to_type_layout(&type_tag) else {
1027            debug_assert!(false);
1028            continue;
1029        };
1030
1031        let Some(annotated_layout) = context.type_tag_to_annotated_type_layout(&type_tag) else {
1032            debug_assert!(false);
1033            continue;
1034        };
1035
1036        let blob = safe_unwrap!(value.borrow().typed_serialize(&layout));
1037        safe_unwrap!(MoveValue::visit_deserialize(
1038            &blob,
1039            &annotated_layout,
1040            &mut Traversal {
1041                state: LookingFor::Wrapped,
1042                ids,
1043                uid: &uid,
1044            },
1045        ));
1046    }
1047    Ok(())
1048}