Skip to main content

iota_move_natives_latest/
test_scenario.rs

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