Skip to main content

sui_core/authority/
shared_object_version_manager.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use mysten_common::ZipDebugEqIteratorExt;
5use mysten_common::debug_fatal;
6
7use crate::authority::AuthorityPerEpochStore;
8use crate::authority::authority_per_epoch_store::CancelConsensusCertificateReason;
9use crate::execution_cache::ObjectCacheRead;
10use either::Either;
11use std::collections::BTreeMap;
12use std::collections::HashMap;
13use std::collections::HashSet;
14use sui_types::SUI_ACCUMULATOR_ROOT_OBJECT_ID;
15use sui_types::SUI_CLOCK_OBJECT_ID;
16use sui_types::SUI_CLOCK_OBJECT_SHARED_VERSION;
17use sui_types::base_types::ConsensusObjectSequenceKey;
18use sui_types::base_types::ConsensusObjectVersion;
19use sui_types::base_types::ObjectID;
20use sui_types::base_types::SystemObjectVersions;
21use sui_types::base_types::TransactionDigest;
22use sui_types::committee::EpochId;
23use sui_types::crypto::RandomnessRound;
24use sui_types::effects::{TransactionEffects, TransactionEffectsAPI};
25use sui_types::executable_transaction::VerifiedExecutableTransaction;
26use sui_types::executable_transaction::VerifiedExecutableTransactionWithAliases;
27use sui_types::storage::{
28    ObjectKey, transaction_non_shared_input_object_keys, transaction_receiving_object_keys,
29};
30use sui_types::transaction::SharedObjectMutability;
31use sui_types::transaction::{SharedInputObject, TransactionDataAPI, TransactionKey};
32use sui_types::{
33    IMPLICITLY_READ_SYSTEM_OBJECTS, SUI_RANDOMNESS_STATE_OBJECT_ID, base_types::SequenceNumber,
34    error::SuiResult,
35};
36use tracing::trace;
37
38pub struct SharedObjVerManager {}
39
40/// Version assignments for a single transaction
41#[derive(Debug, Clone, PartialEq, Eq)]
42pub struct AssignedVersions {
43    pub shared_object_versions: Vec<(ConsensusObjectSequenceKey, SequenceNumber)>,
44    /// Versions of system objects, keyed by object ID, that this transaction may read during
45    /// execution but that are not part of its declared shared inputs. Each version is assigned
46    /// deterministically during consensus sequencing, so that every validator reads the same
47    /// version of the object.
48    ///
49    /// Today this holds at most the accumulator root version (as of the beginning of the consensus
50    /// commit this transaction belongs to). The accumulator root qualifies because it is written at
51    /// the end of every commit, so there is always a well-defined prior version to read from. More
52    /// system objects will be added over time.
53    pub system_object_versions: SystemObjectVersions,
54}
55
56impl AssignedVersions {
57    pub fn new(
58        shared_object_versions: Vec<(ConsensusObjectSequenceKey, SequenceNumber)>,
59        system_object_versions: SystemObjectVersions,
60    ) -> Self {
61        Self {
62            shared_object_versions,
63            system_object_versions,
64        }
65    }
66
67    pub fn empty() -> Self {
68        Self::new(vec![], SystemObjectVersions::empty())
69    }
70
71    /// Construct with only the accumulator root as the system object read during execution. The
72    /// accumulator root is the sole such object today; production callers build the full
73    /// `system_object_versions` map directly.
74    #[cfg(test)]
75    pub fn new_for_testing(
76        shared_object_versions: Vec<(ConsensusObjectSequenceKey, SequenceNumber)>,
77        accumulator_version: Option<SequenceNumber>,
78    ) -> Self {
79        Self::new(
80            shared_object_versions,
81            SystemObjectVersions::new(accumulator_version.map(|v| ConsensusObjectVersion {
82                initial_shared_version: sui_types::object::OBJECT_START_VERSION,
83                version: v,
84            })),
85        )
86    }
87
88    /// The accumulator root version this transaction reads, if any.
89    pub fn accumulator_version(&self) -> Option<SequenceNumber> {
90        self.system_object_versions
91            .get(&SUI_ACCUMULATOR_ROOT_OBJECT_ID)
92            .map(|v| v.version)
93    }
94
95    pub fn iter(&self) -> impl Iterator<Item = &(ConsensusObjectSequenceKey, SequenceNumber)> {
96        self.shared_object_versions.iter()
97    }
98
99    pub fn as_slice(&self) -> &[(ConsensusObjectSequenceKey, SequenceNumber)] {
100        &self.shared_object_versions
101    }
102}
103
104#[derive(Default, Debug, PartialEq, Eq)]
105pub struct AssignedTxAndVersions(pub Vec<(TransactionKey, AssignedVersions)>);
106
107impl AssignedTxAndVersions {
108    pub fn new(assigned_versions: Vec<(TransactionKey, AssignedVersions)>) -> Self {
109        Self(assigned_versions)
110    }
111
112    pub fn into_map(self) -> HashMap<TransactionKey, AssignedVersions> {
113        self.0.into_iter().collect()
114    }
115}
116
117/// A wrapper around things that can be scheduled for execution by the assigning of
118/// shared object versions.
119#[derive(Clone)]
120pub enum Schedulable<T = VerifiedExecutableTransaction> {
121    Transaction(T),
122    RandomnessStateUpdate(EpochId, RandomnessRound),
123    AccumulatorSettlement(EpochId, u64 /* checkpoint height */),
124    ConsensusCommitPrologue(EpochId, u64 /* round */, u32 /* sub_dag_index */),
125}
126
127impl From<VerifiedExecutableTransaction> for Schedulable<VerifiedExecutableTransaction> {
128    fn from(tx: VerifiedExecutableTransaction) -> Self {
129        Schedulable::Transaction(tx)
130    }
131}
132
133impl From<Schedulable<VerifiedExecutableTransactionWithAliases>>
134    for Schedulable<VerifiedExecutableTransaction>
135{
136    fn from(schedulable: Schedulable<VerifiedExecutableTransactionWithAliases>) -> Self {
137        match schedulable {
138            Schedulable::Transaction(tx) => Schedulable::Transaction(tx.into_tx()),
139            Schedulable::RandomnessStateUpdate(epoch, round) => {
140                Schedulable::RandomnessStateUpdate(epoch, round)
141            }
142            Schedulable::AccumulatorSettlement(epoch, checkpoint_height) => {
143                Schedulable::AccumulatorSettlement(epoch, checkpoint_height)
144            }
145            Schedulable::ConsensusCommitPrologue(epoch, round, sub_dag_index) => {
146                Schedulable::ConsensusCommitPrologue(epoch, round, sub_dag_index)
147            }
148        }
149    }
150}
151
152// AsTx is like Deref, in that it allows us to use either refs or values in Schedulable.
153// Deref does not work because it conflicts with the impl of Deref for VerifiedExecutableTransaction.
154pub trait AsTx {
155    fn as_tx(&self) -> &VerifiedExecutableTransaction;
156}
157
158impl AsTx for VerifiedExecutableTransaction {
159    fn as_tx(&self) -> &VerifiedExecutableTransaction {
160        self
161    }
162}
163
164impl AsTx for &'_ VerifiedExecutableTransaction {
165    fn as_tx(&self) -> &VerifiedExecutableTransaction {
166        self
167    }
168}
169
170impl AsTx for VerifiedExecutableTransactionWithAliases {
171    fn as_tx(&self) -> &VerifiedExecutableTransaction {
172        self.tx()
173    }
174}
175
176impl AsTx for &'_ VerifiedExecutableTransactionWithAliases {
177    fn as_tx(&self) -> &VerifiedExecutableTransaction {
178        self.tx()
179    }
180}
181
182impl Schedulable<&'_ VerifiedExecutableTransaction> {
183    // Cannot use the blanket ToOwned trait impl because it just calls clone.
184    pub fn to_owned_schedulable(&self) -> Schedulable<VerifiedExecutableTransaction> {
185        match self {
186            Schedulable::Transaction(tx) => Schedulable::Transaction((*tx).clone()),
187            Schedulable::RandomnessStateUpdate(epoch, round) => {
188                Schedulable::RandomnessStateUpdate(*epoch, *round)
189            }
190            Schedulable::AccumulatorSettlement(epoch, checkpoint_height) => {
191                Schedulable::AccumulatorSettlement(*epoch, *checkpoint_height)
192            }
193            Schedulable::ConsensusCommitPrologue(epoch, round, sub_dag_index) => {
194                Schedulable::ConsensusCommitPrologue(*epoch, *round, *sub_dag_index)
195            }
196        }
197    }
198}
199
200impl<T> Schedulable<T> {
201    pub fn as_tx(&self) -> Option<&VerifiedExecutableTransaction>
202    where
203        T: AsTx,
204    {
205        match self {
206            Schedulable::Transaction(tx) => Some(tx.as_tx()),
207            Schedulable::RandomnessStateUpdate(_, _) => None,
208            Schedulable::AccumulatorSettlement(_, _) => None,
209            Schedulable::ConsensusCommitPrologue(_, _, _) => None,
210        }
211    }
212
213    pub fn shared_input_objects(
214        &self,
215        epoch_store: &AuthorityPerEpochStore,
216    ) -> impl Iterator<Item = SharedInputObject> + '_
217    where
218        T: AsTx,
219    {
220        match self {
221            Schedulable::Transaction(tx) => Either::Left(tx.as_tx().shared_input_objects()),
222            Schedulable::RandomnessStateUpdate(_, _) => {
223                Either::Right(std::iter::once(SharedInputObject {
224                    id: SUI_RANDOMNESS_STATE_OBJECT_ID,
225                    initial_shared_version: epoch_store
226                        .epoch_start_config()
227                        .randomness_obj_initial_shared_version()
228                        .expect("randomness obj initial shared version should be set"),
229                    mutability: SharedObjectMutability::Mutable,
230                }))
231            }
232            Schedulable::AccumulatorSettlement(_, _) => {
233                Either::Right(std::iter::once(SharedInputObject {
234                    id: SUI_ACCUMULATOR_ROOT_OBJECT_ID,
235                    initial_shared_version: epoch_store
236                        .epoch_start_config()
237                        .accumulator_root_obj_initial_shared_version()
238                        .expect("accumulator root obj initial shared version should be set"),
239                    mutability: SharedObjectMutability::Mutable,
240                }))
241            }
242            Schedulable::ConsensusCommitPrologue(_, _, _) => {
243                Either::Right(std::iter::once(SharedInputObject {
244                    id: SUI_CLOCK_OBJECT_ID,
245                    initial_shared_version: SUI_CLOCK_OBJECT_SHARED_VERSION,
246                    mutability: SharedObjectMutability::Mutable,
247                }))
248            }
249        }
250    }
251
252    pub fn non_shared_input_object_keys(&self) -> Vec<ObjectKey>
253    where
254        T: AsTx,
255    {
256        match self {
257            Schedulable::Transaction(tx) => transaction_non_shared_input_object_keys(tx.as_tx())
258                .expect("Transaction input should have been verified"),
259            Schedulable::RandomnessStateUpdate(_, _) => vec![],
260            Schedulable::AccumulatorSettlement(_, _) => vec![],
261            Schedulable::ConsensusCommitPrologue(_, _, _) => vec![],
262        }
263    }
264
265    pub fn receiving_object_keys(&self) -> Vec<ObjectKey>
266    where
267        T: AsTx,
268    {
269        match self {
270            Schedulable::Transaction(tx) => transaction_receiving_object_keys(tx.as_tx()),
271            Schedulable::RandomnessStateUpdate(_, _) => vec![],
272            Schedulable::AccumulatorSettlement(_, _) => vec![],
273            Schedulable::ConsensusCommitPrologue(_, _, _) => vec![],
274        }
275    }
276
277    pub fn key(&self) -> TransactionKey
278    where
279        T: AsTx,
280    {
281        match self {
282            Schedulable::Transaction(tx) => tx.as_tx().key(),
283            Schedulable::RandomnessStateUpdate(epoch, round) => {
284                TransactionKey::RandomnessRound(*epoch, *round)
285            }
286            Schedulable::AccumulatorSettlement(epoch, checkpoint_height) => {
287                TransactionKey::AccumulatorSettlement(*epoch, *checkpoint_height)
288            }
289            Schedulable::ConsensusCommitPrologue(epoch, round, sub_dag_index) => {
290                TransactionKey::ConsensusCommitPrologue(*epoch, *round, *sub_dag_index)
291            }
292        }
293    }
294}
295
296#[must_use]
297#[derive(Default, Eq, PartialEq, Debug)]
298pub struct ConsensusSharedObjVerAssignment {
299    pub shared_input_next_versions: HashMap<ConsensusObjectSequenceKey, SequenceNumber>,
300    pub assigned_versions: AssignedTxAndVersions,
301}
302
303impl SharedObjVerManager {
304    pub fn assign_versions_from_consensus<'a, T>(
305        epoch_store: &AuthorityPerEpochStore,
306        cache_reader: &dyn ObjectCacheRead,
307        assignables: impl Iterator<Item = &'a Schedulable<T>> + Clone,
308        cancelled_txns: &BTreeMap<TransactionDigest, CancelConsensusCertificateReason>,
309    ) -> SuiResult<ConsensusSharedObjVerAssignment>
310    where
311        T: AsTx + 'a,
312    {
313        let mut shared_input_next_versions = get_or_init_versions(
314            assignables
315                .clone()
316                .flat_map(|a| a.shared_input_objects(epoch_store)),
317            epoch_store,
318            cache_reader,
319        )?;
320        let mut assigned_versions = Vec::new();
321        for assignable in assignables {
322            assert!(
323                !matches!(assignable, Schedulable::AccumulatorSettlement(_, _))
324                    || epoch_store.accumulators_enabled(),
325                "AccumulatorSettlement should not be scheduled when accumulators are disabled"
326            );
327
328            let cert_assigned_versions = Self::assign_versions_for_certificate(
329                epoch_store,
330                assignable,
331                &mut shared_input_next_versions,
332                cancelled_txns,
333            );
334            assigned_versions.push((assignable.key(), cert_assigned_versions));
335        }
336
337        Ok(ConsensusSharedObjVerAssignment {
338            shared_input_next_versions,
339            assigned_versions: AssignedTxAndVersions::new(assigned_versions),
340        })
341    }
342
343    pub fn assign_versions_from_effects(
344        certs_and_effects: &[(
345            &VerifiedExecutableTransaction,
346            &TransactionEffects,
347            // Accumulator version
348            Option<SequenceNumber>,
349        )],
350        epoch_store: &AuthorityPerEpochStore,
351        cache_reader: &dyn ObjectCacheRead,
352    ) -> AssignedTxAndVersions {
353        // We don't care about the results since we can use effects to assign versions.
354        // But we must call it to make sure whenever a consensus object is touched the first time
355        // during an epoch, either through consensus or through checkpoint executor,
356        // its next version must be initialized. This is because we initialize the next version
357        // of a consensus object in an epoch by reading the current version from the object store.
358        // This must be done before we mutate it the first time, otherwise we would be initializing
359        // it with the wrong version.
360        let _ = get_or_init_versions(
361            certs_and_effects.iter().flat_map(|(cert, _, _)| {
362                cert.transaction_data().shared_input_objects().into_iter()
363            }),
364            epoch_store,
365            cache_reader,
366        );
367        let mut assigned_versions = Vec::new();
368        for (cert, effects, accumulator_version) in certs_and_effects {
369            let initial_version_map: BTreeMap<_, _> = cert
370                .transaction_data()
371                .shared_input_objects()
372                .into_iter()
373                .map(|input| input.into_id_and_version())
374                .collect();
375            // When we add more implicitly read system objects, retrieve their accessed versions
376            // from this map and pass them to `SystemObjectVersions`.
377            let accessed_versions: BTreeMap<ObjectID, SequenceNumber> = effects
378                .accessed_consensus_objects()
379                .into_iter()
380                .map(|iso| iso.id_and_version())
381                .collect();
382            let cert_assigned_versions: Vec<_> = accessed_versions
383                .iter()
384                .filter_map(|(id, version)| {
385                    let v = initial_version_map
386                        .get(id)
387                        .map(|initial_version| ((*id, *initial_version), *version));
388                    if v.is_none() {
389                        debug_assert!(
390                            IMPLICITLY_READ_SYSTEM_OBJECTS.contains(id),
391                            "accessed consensus object is neither a declared input nor a known implicitly read system object: \
392                             accessed={accessed_versions:?} declared={initial_version_map:?}"
393                        );
394                    }
395                    v
396                })
397                .collect();
398            if let (Some(effects_version), Some(sequenced_version)) = (
399                accessed_versions.get(&SUI_ACCUMULATOR_ROOT_OBJECT_ID),
400                accumulator_version,
401            ) && !effects_version.is_cancelled()
402                && effects_version != sequenced_version
403            {
404                debug_fatal!(
405                    "accumulator root version from effects {:?} disagrees \
406                        with the reconstructed accumulator version {:?} for tx {:?}",
407                    effects_version,
408                    sequenced_version,
409                    cert.digest()
410                );
411            }
412            // Note that for accumulator version, we cannot rely on the one from effects yet, since it won't
413            // be produced until implicitly read system objects are fully shipped. But the old object funds withdraw
414            // still need it. Hence we always use the one provided from the caller (i.e. checkpoint executor).
415            let system_object_versions =
416                SystemObjectVersions::new(accumulator_version.map(|version| {
417                    let initial_shared_version = epoch_store
418                        .epoch_start_config()
419                        .accumulator_root_obj_initial_shared_version()
420                        .expect(
421                            "initial shared version must be known for an implicitly read system object",
422                        );
423                    ConsensusObjectVersion {
424                        initial_shared_version,
425                        version,
426                    }
427                }));
428            let tx_key = cert.key();
429            trace!(
430                ?tx_key,
431                ?cert_assigned_versions,
432                ?system_object_versions,
433                "assigned consensus object versions from effects"
434            );
435            assigned_versions.push((
436                tx_key,
437                AssignedVersions::new(cert_assigned_versions, system_object_versions),
438            ));
439        }
440        AssignedTxAndVersions::new(assigned_versions)
441    }
442
443    pub fn assign_versions_for_certificate(
444        epoch_store: &AuthorityPerEpochStore,
445        assignable: &Schedulable<impl AsTx>,
446        shared_input_next_versions: &mut HashMap<ConsensusObjectSequenceKey, SequenceNumber>,
447        cancelled_txns: &BTreeMap<TransactionDigest, CancelConsensusCertificateReason>,
448    ) -> AssignedVersions {
449        let shared_input_objects: Vec<_> = assignable.shared_input_objects(epoch_store).collect();
450
451        let accumulator_version = if epoch_store.accumulators_enabled() {
452            let accumulator_initial_version = epoch_store
453                .epoch_start_config()
454                .accumulator_root_obj_initial_shared_version()
455                .expect("accumulator root obj initial shared version should be set when accumulators are enabled");
456
457            let accumulator_version = *shared_input_next_versions
458                .get(&(SUI_ACCUMULATOR_ROOT_OBJECT_ID, accumulator_initial_version))
459                .expect("accumulator object must be in shared_input_next_versions when withdraws are enabled");
460
461            Some(ConsensusObjectVersion {
462                initial_shared_version: accumulator_initial_version,
463                version: accumulator_version,
464            })
465        } else {
466            None
467        };
468        let system_object_versions = SystemObjectVersions::new(accumulator_version);
469
470        if shared_input_objects.is_empty() {
471            // No shared object used by this transaction. No need to assign versions.
472            return AssignedVersions::new(vec![], system_object_versions);
473        }
474
475        let tx_key = assignable.key();
476
477        // Check if the transaction is cancelled due to congestion.
478        let cancellation_info = tx_key
479            .as_digest()
480            .and_then(|tx_digest| cancelled_txns.get(tx_digest));
481        let congested_objects_info: Option<HashSet<_>> =
482            if let Some(CancelConsensusCertificateReason::CongestionOnObjects(congested_objects)) =
483                &cancellation_info
484            {
485                Some(congested_objects.iter().cloned().collect())
486            } else {
487                None
488            };
489        let txn_cancelled = cancellation_info.is_some();
490
491        let mut input_object_keys = assignable.non_shared_input_object_keys();
492        let mut assigned_versions = Vec::with_capacity(shared_input_objects.len());
493        let mut is_exclusively_accessed_input = Vec::with_capacity(shared_input_objects.len());
494        // Record receiving object versions towards the shared version computation.
495        let receiving_object_keys = assignable.receiving_object_keys();
496        input_object_keys.extend(receiving_object_keys);
497
498        if txn_cancelled {
499            // For cancelled transaction due to congestion, assign special versions to all shared objects.
500            // Note that new lamport version does not depend on any shared objects.
501            for SharedInputObject {
502                id,
503                initial_shared_version,
504                ..
505            } in shared_input_objects.iter()
506            {
507                let assigned_version = match cancellation_info {
508                    Some(CancelConsensusCertificateReason::CongestionOnObjects(_)) => {
509                        if congested_objects_info
510                            .as_ref()
511                            .is_some_and(|info| info.contains(id))
512                        {
513                            SequenceNumber::CONGESTED
514                        } else {
515                            SequenceNumber::CANCELLED_READ
516                        }
517                    }
518                    Some(CancelConsensusCertificateReason::DkgFailed) => {
519                        if id == &SUI_RANDOMNESS_STATE_OBJECT_ID {
520                            SequenceNumber::RANDOMNESS_UNAVAILABLE
521                        } else {
522                            SequenceNumber::CANCELLED_READ
523                        }
524                    }
525                    None => unreachable!("cancelled transaction should have cancellation info"),
526                };
527                assigned_versions.push(((*id, *initial_shared_version), assigned_version));
528                is_exclusively_accessed_input.push(false);
529            }
530        } else {
531            for (
532                SharedInputObject {
533                    id,
534                    initial_shared_version,
535                    mutability,
536                },
537                assigned_version,
538            ) in shared_input_objects.iter().map(|obj| {
539                (
540                    obj,
541                    *shared_input_next_versions
542                        .get(&obj.id_and_version())
543                        .unwrap(),
544                )
545            }) {
546                assigned_versions.push(((*id, *initial_shared_version), assigned_version));
547                input_object_keys.push(ObjectKey(*id, assigned_version));
548                is_exclusively_accessed_input.push(mutability.is_exclusive());
549            }
550        }
551
552        let next_version =
553            SequenceNumber::lamport_increment(input_object_keys.iter().map(|obj| obj.1));
554        assert!(
555            next_version.is_valid(),
556            "Assigned version must be valid. Got {:?}",
557            next_version
558        );
559
560        if !txn_cancelled {
561            // Update the next version for the shared objects.
562            assigned_versions
563                .iter()
564                .zip_debug_eq(is_exclusively_accessed_input)
565                .filter_map(|((id, _), mutable)| {
566                    if mutable {
567                        Some((*id, next_version))
568                    } else {
569                        None
570                    }
571                })
572                .for_each(|(id, version)| {
573                    assert!(
574                        version.is_valid(),
575                        "Assigned version must be a valid version."
576                    );
577                    shared_input_next_versions
578                        .insert(id, version)
579                        .expect("Object must exist in shared_input_next_versions.");
580                });
581        }
582
583        trace!(
584            ?tx_key,
585            ?assigned_versions,
586            ?next_version,
587            ?txn_cancelled,
588            "locking shared objects"
589        );
590
591        AssignedVersions::new(assigned_versions, system_object_versions)
592    }
593}
594
595fn get_or_init_versions<'a>(
596    shared_input_objects: impl Iterator<Item = SharedInputObject> + 'a,
597    epoch_store: &AuthorityPerEpochStore,
598    cache_reader: &dyn ObjectCacheRead,
599) -> SuiResult<HashMap<ConsensusObjectSequenceKey, SequenceNumber>> {
600    let mut shared_input_objects: Vec<_> = shared_input_objects
601        .map(|so| so.into_id_and_version())
602        .collect();
603
604    if epoch_store.accumulators_enabled() {
605        shared_input_objects.push((
606            SUI_ACCUMULATOR_ROOT_OBJECT_ID,
607            epoch_store
608                .epoch_start_config()
609                .accumulator_root_obj_initial_shared_version()
610                .expect("accumulator root obj initial shared version should be set"),
611        ));
612    }
613
614    shared_input_objects.sort();
615    shared_input_objects.dedup();
616
617    epoch_store.get_or_init_next_object_versions(&shared_input_objects, cache_reader)
618}
619
620#[cfg(test)]
621mod tests {
622    use super::*;
623
624    use crate::authority::AuthorityState;
625    use crate::authority::shared_object_version_manager::{
626        ConsensusSharedObjVerAssignment, SharedObjVerManager,
627    };
628    use crate::authority::test_authority_builder::TestAuthorityBuilder;
629    use std::collections::{BTreeMap, HashMap};
630    use std::sync::Arc;
631    use sui_protocol_config::ProtocolConfig;
632    use sui_test_transaction_builder::TestTransactionBuilder;
633    use sui_types::base_types::{ObjectID, SequenceNumber, SuiAddress};
634    use sui_types::crypto::{RandomnessRound, get_account_key_pair};
635    use sui_types::digests::ObjectDigest;
636    use sui_types::effects::TestEffectsBuilder;
637    use sui_types::executable_transaction::{
638        CertificateProof, ExecutableTransaction, VerifiedExecutableTransaction,
639    };
640
641    use sui_types::object::Object;
642    use sui_types::transaction::{ObjectArg, SenderSignedData, VerifiedTransaction};
643
644    use sui_types::gas_coin::GAS;
645    use sui_types::transaction::FundsWithdrawalArg;
646    use sui_types::{SUI_ACCUMULATOR_ROOT_OBJECT_ID, SUI_RANDOMNESS_STATE_OBJECT_ID};
647
648    #[tokio::test]
649    async fn test_assign_versions_from_consensus_basic() {
650        let shared_object = Object::shared_for_testing();
651        let id = shared_object.id();
652        let init_shared_version = shared_object.owner.start_version().unwrap();
653        let authority = TestAuthorityBuilder::new()
654            .with_starting_objects(std::slice::from_ref(&shared_object))
655            .build()
656            .await;
657        let certs = [
658            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, true)], 3),
659            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, false)], 5),
660            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, true)], 9),
661            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, true)], 11),
662        ];
663        let epoch_store = authority.epoch_store_for_testing();
664        let assignables = certs
665            .iter()
666            .map(Schedulable::Transaction)
667            .collect::<Vec<_>>();
668        let ConsensusSharedObjVerAssignment {
669            shared_input_next_versions,
670            assigned_versions,
671        } = SharedObjVerManager::assign_versions_from_consensus(
672            &epoch_store,
673            authority.get_object_cache_reader().as_ref(),
674            assignables.iter(),
675            &BTreeMap::new(),
676        )
677        .unwrap();
678        // Check that the shared object's next version is always initialized in the epoch store.
679        assert_eq!(
680            epoch_store
681                .get_next_object_version(&id, init_shared_version)
682                .unwrap(),
683            init_shared_version
684        );
685        // Check that the final version of the shared object is the lamport version of the last
686        // transaction.
687        assert_eq!(
688            *shared_input_next_versions
689                .get(&(id, init_shared_version))
690                .unwrap(),
691            SequenceNumber::from_u64(12)
692        );
693        // Check that the version assignment for each transaction is correct.
694        // For a transaction that uses the shared object with mutable=false, it won't update the version
695        // using lamport version, hence the next transaction will use the same version number.
696        // In the following case, certs[2] has the same assignment as certs[1] for this reason.
697        let expected_accumulator_version = SequenceNumber::from_u64(1);
698        assert_eq!(
699            assigned_versions.0,
700            vec![
701                (
702                    certs[0].key(),
703                    AssignedVersions::new_for_testing(
704                        vec![((id, init_shared_version), init_shared_version)],
705                        Some(expected_accumulator_version)
706                    )
707                ),
708                (
709                    certs[1].key(),
710                    AssignedVersions::new_for_testing(
711                        vec![((id, init_shared_version), SequenceNumber::from_u64(4))],
712                        Some(expected_accumulator_version)
713                    )
714                ),
715                (
716                    certs[2].key(),
717                    AssignedVersions::new_for_testing(
718                        vec![((id, init_shared_version), SequenceNumber::from_u64(4))],
719                        Some(expected_accumulator_version)
720                    )
721                ),
722                (
723                    certs[3].key(),
724                    AssignedVersions::new_for_testing(
725                        vec![((id, init_shared_version), SequenceNumber::from_u64(10))],
726                        Some(expected_accumulator_version)
727                    )
728                ),
729            ]
730        );
731    }
732
733    #[tokio::test]
734    async fn test_assign_versions_from_consensus_with_randomness() {
735        let authority = TestAuthorityBuilder::new().build().await;
736        let epoch_store = authority.epoch_store_for_testing();
737        let randomness_obj_version = epoch_store
738            .epoch_start_config()
739            .randomness_obj_initial_shared_version()
740            .unwrap();
741        let certs = [
742            VerifiedExecutableTransaction::new_system(
743                VerifiedTransaction::new_randomness_state_update(
744                    epoch_store.epoch(),
745                    RandomnessRound::new(1),
746                    vec![],
747                    randomness_obj_version,
748                ),
749                epoch_store.epoch(),
750            ),
751            generate_shared_objs_tx_with_gas_version(
752                &[(
753                    SUI_RANDOMNESS_STATE_OBJECT_ID,
754                    randomness_obj_version,
755                    // This can only be false since it's not allowed to use randomness object with mutable=true.
756                    false,
757                )],
758                3,
759            ),
760            generate_shared_objs_tx_with_gas_version(
761                &[(
762                    SUI_RANDOMNESS_STATE_OBJECT_ID,
763                    randomness_obj_version,
764                    false,
765                )],
766                5,
767            ),
768        ];
769        let assignables = certs
770            .iter()
771            .map(Schedulable::Transaction)
772            .collect::<Vec<_>>();
773        let ConsensusSharedObjVerAssignment {
774            shared_input_next_versions,
775            assigned_versions,
776        } = SharedObjVerManager::assign_versions_from_consensus(
777            &epoch_store,
778            authority.get_object_cache_reader().as_ref(),
779            assignables.iter(),
780            &BTreeMap::new(),
781        )
782        .unwrap();
783        // Check that the randomness object's next version is initialized.
784        assert_eq!(
785            epoch_store
786                .get_next_object_version(&SUI_RANDOMNESS_STATE_OBJECT_ID, randomness_obj_version)
787                .unwrap(),
788            randomness_obj_version
789        );
790        let next_randomness_obj_version = randomness_obj_version.next();
791        assert_eq!(
792            *shared_input_next_versions
793                .get(&(SUI_RANDOMNESS_STATE_OBJECT_ID, randomness_obj_version))
794                .unwrap(),
795            // Randomness object's version is only incremented by 1 regardless of lamport version.
796            next_randomness_obj_version
797        );
798        let expected_accumulator_version = SequenceNumber::from_u64(1);
799        assert_eq!(
800            assigned_versions.0,
801            vec![
802                (
803                    certs[0].key(),
804                    AssignedVersions::new_for_testing(
805                        vec![(
806                            (SUI_RANDOMNESS_STATE_OBJECT_ID, randomness_obj_version),
807                            randomness_obj_version
808                        )],
809                        Some(expected_accumulator_version)
810                    )
811                ),
812                (
813                    certs[1].key(),
814                    // It is critical that the randomness object version is updated before the assignment.
815                    AssignedVersions::new_for_testing(
816                        vec![(
817                            (SUI_RANDOMNESS_STATE_OBJECT_ID, randomness_obj_version),
818                            next_randomness_obj_version
819                        )],
820                        Some(expected_accumulator_version)
821                    )
822                ),
823                (
824                    certs[2].key(),
825                    // It is critical that the randomness object version is updated before the assignment.
826                    AssignedVersions::new_for_testing(
827                        vec![(
828                            (SUI_RANDOMNESS_STATE_OBJECT_ID, randomness_obj_version),
829                            next_randomness_obj_version
830                        )],
831                        Some(expected_accumulator_version)
832                    )
833                ),
834            ]
835        );
836    }
837
838    // Tests shared object version assignment for cancelled transaction.
839    #[tokio::test]
840    async fn test_assign_versions_from_consensus_with_cancellation() {
841        let shared_object_1 = Object::shared_for_testing();
842        let shared_object_2 = Object::shared_for_testing();
843        let id1 = shared_object_1.id();
844        let id2 = shared_object_2.id();
845        let init_shared_version_1 = shared_object_1.owner.start_version().unwrap();
846        let init_shared_version_2 = shared_object_2.owner.start_version().unwrap();
847        let authority = TestAuthorityBuilder::new()
848            .with_starting_objects(&[shared_object_1.clone(), shared_object_2.clone()])
849            .build()
850            .await;
851        let randomness_obj_version = authority
852            .epoch_store_for_testing()
853            .epoch_start_config()
854            .randomness_obj_initial_shared_version()
855            .unwrap();
856
857        // Generate 5 transactions for testing.
858        //   tx1: shared_object_1, shared_object_2, owned_object_version = 3
859        //   tx2: shared_object_1, shared_object_2, owned_object_version = 5
860        //   tx3: shared_object_1, owned_object_version = 1
861        //   tx4: shared_object_1, shared_object_2, owned_object_version = 9
862        //   tx5: shared_object_1, shared_object_2, owned_object_version = 11
863        //
864        // Later, we cancel transaction 2 and 4 due to congestion, and 5 due to DKG failure.
865        // Expected outcome:
866        //   tx1: both shared objects assign version 1, lamport version = 4
867        //   tx2: shared objects assign cancelled version, lamport version = 6 due to gas object version = 5
868        //   tx3: shared object 1 assign version 4, lamport version = 5
869        //   tx4: shared objects assign cancelled version, lamport version = 10 due to gas object version = 9
870        //   tx5: shared objects assign cancelled version, lamport version = 12 due to gas object version = 11
871        let certs = [
872            generate_shared_objs_tx_with_gas_version(
873                &[
874                    (id1, init_shared_version_1, true),
875                    (id2, init_shared_version_2, true),
876                ],
877                3,
878            ),
879            generate_shared_objs_tx_with_gas_version(
880                &[
881                    (id1, init_shared_version_1, true),
882                    (id2, init_shared_version_2, true),
883                ],
884                5,
885            ),
886            generate_shared_objs_tx_with_gas_version(&[(id1, init_shared_version_1, true)], 1),
887            generate_shared_objs_tx_with_gas_version(
888                &[
889                    (id1, init_shared_version_1, true),
890                    (id2, init_shared_version_2, true),
891                ],
892                9,
893            ),
894            generate_shared_objs_tx_with_gas_version(
895                &[
896                    (
897                        SUI_RANDOMNESS_STATE_OBJECT_ID,
898                        randomness_obj_version,
899                        false,
900                    ),
901                    (id2, init_shared_version_2, true),
902                ],
903                11,
904            ),
905        ];
906        let epoch_store = authority.epoch_store_for_testing();
907
908        // Cancel transactions 2 and 4 due to congestion.
909        let cancelled_txns: BTreeMap<TransactionDigest, CancelConsensusCertificateReason> = [
910            (
911                *certs[1].digest(),
912                CancelConsensusCertificateReason::CongestionOnObjects(vec![id1]),
913            ),
914            (
915                *certs[3].digest(),
916                CancelConsensusCertificateReason::CongestionOnObjects(vec![id2]),
917            ),
918            (
919                *certs[4].digest(),
920                CancelConsensusCertificateReason::DkgFailed,
921            ),
922        ]
923        .into_iter()
924        .collect();
925
926        let assignables = certs
927            .iter()
928            .map(Schedulable::Transaction)
929            .collect::<Vec<_>>();
930
931        // Run version assignment logic.
932        let ConsensusSharedObjVerAssignment {
933            mut shared_input_next_versions,
934            assigned_versions,
935        } = SharedObjVerManager::assign_versions_from_consensus(
936            &epoch_store,
937            authority.get_object_cache_reader().as_ref(),
938            assignables.iter(),
939            &cancelled_txns,
940        )
941        .unwrap();
942
943        // Check that the final version of the shared object is the lamport version of the last
944        // transaction.
945        shared_input_next_versions
946            .remove(&(SUI_ACCUMULATOR_ROOT_OBJECT_ID, SequenceNumber::from_u64(1)));
947        assert_eq!(
948            shared_input_next_versions,
949            HashMap::from([
950                ((id1, init_shared_version_1), SequenceNumber::from_u64(5)), // determined by tx3
951                ((id2, init_shared_version_2), SequenceNumber::from_u64(4)), // determined by tx1
952                (
953                    (SUI_RANDOMNESS_STATE_OBJECT_ID, randomness_obj_version),
954                    SequenceNumber::from_u64(1)
955                ), // not mutable
956            ])
957        );
958
959        // Check that the version assignment for each transaction is correct.
960        let expected_accumulator_version = SequenceNumber::from_u64(1);
961        assert_eq!(
962            assigned_versions.0,
963            vec![
964                (
965                    certs[0].key(),
966                    AssignedVersions::new_for_testing(
967                        vec![
968                            ((id1, init_shared_version_1), init_shared_version_1),
969                            ((id2, init_shared_version_2), init_shared_version_2)
970                        ],
971                        Some(expected_accumulator_version)
972                    )
973                ),
974                (
975                    certs[1].key(),
976                    AssignedVersions::new_for_testing(
977                        vec![
978                            ((id1, init_shared_version_1), SequenceNumber::CONGESTED),
979                            ((id2, init_shared_version_2), SequenceNumber::CANCELLED_READ),
980                        ],
981                        Some(expected_accumulator_version)
982                    )
983                ),
984                (
985                    certs[2].key(),
986                    AssignedVersions::new_for_testing(
987                        vec![((id1, init_shared_version_1), SequenceNumber::from_u64(4))],
988                        Some(expected_accumulator_version)
989                    )
990                ),
991                (
992                    certs[3].key(),
993                    AssignedVersions::new_for_testing(
994                        vec![
995                            ((id1, init_shared_version_1), SequenceNumber::CANCELLED_READ),
996                            ((id2, init_shared_version_2), SequenceNumber::CONGESTED)
997                        ],
998                        Some(expected_accumulator_version)
999                    )
1000                ),
1001                (
1002                    certs[4].key(),
1003                    AssignedVersions::new_for_testing(
1004                        vec![
1005                            (
1006                                (SUI_RANDOMNESS_STATE_OBJECT_ID, randomness_obj_version),
1007                                SequenceNumber::RANDOMNESS_UNAVAILABLE
1008                            ),
1009                            ((id2, init_shared_version_2), SequenceNumber::CANCELLED_READ)
1010                        ],
1011                        Some(expected_accumulator_version)
1012                    )
1013                ),
1014            ]
1015        );
1016    }
1017
1018    #[tokio::test]
1019    async fn test_assign_versions_from_effects() {
1020        let shared_object = Object::shared_for_testing();
1021        let id = shared_object.id();
1022        let init_shared_version = shared_object.owner.start_version().unwrap();
1023        let authority = TestAuthorityBuilder::new()
1024            .with_starting_objects(std::slice::from_ref(&shared_object))
1025            .build()
1026            .await;
1027        let certs = [
1028            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, true)], 3),
1029            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, false)], 5),
1030            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, true)], 9),
1031            generate_shared_objs_tx_with_gas_version(&[(id, init_shared_version, true)], 11),
1032        ];
1033        let effects = [
1034            TestEffectsBuilder::new(certs[0].data()).build(),
1035            TestEffectsBuilder::new(certs[1].data())
1036                .with_shared_input_versions(BTreeMap::from([(id, SequenceNumber::from_u64(4))]))
1037                .build(),
1038            TestEffectsBuilder::new(certs[2].data())
1039                .with_shared_input_versions(BTreeMap::from([(id, SequenceNumber::from_u64(4))]))
1040                .build(),
1041            TestEffectsBuilder::new(certs[3].data())
1042                .with_shared_input_versions(BTreeMap::from([(id, SequenceNumber::from_u64(10))]))
1043                .build(),
1044        ];
1045        let epoch_store = authority.epoch_store_for_testing();
1046        let assigned_versions = SharedObjVerManager::assign_versions_from_effects(
1047            certs
1048                .iter()
1049                .zip_debug_eq(effects.iter())
1050                .map(|(cert, effect)| (cert, effect, None))
1051                .collect::<Vec<_>>()
1052                .as_slice(),
1053            &epoch_store,
1054            authority.get_object_cache_reader().as_ref(),
1055        );
1056        // Check that the shared object's next version is always initialized in the epoch store.
1057        assert_eq!(
1058            epoch_store
1059                .get_next_object_version(&id, init_shared_version)
1060                .unwrap(),
1061            init_shared_version
1062        );
1063        assert_eq!(
1064            assigned_versions.0,
1065            vec![
1066                (
1067                    certs[0].key(),
1068                    AssignedVersions::new_for_testing(
1069                        vec![((id, init_shared_version), init_shared_version)],
1070                        None
1071                    )
1072                ),
1073                (
1074                    certs[1].key(),
1075                    AssignedVersions::new_for_testing(
1076                        vec![((id, init_shared_version), SequenceNumber::from_u64(4))],
1077                        None
1078                    )
1079                ),
1080                (
1081                    certs[2].key(),
1082                    AssignedVersions::new_for_testing(
1083                        vec![((id, init_shared_version), SequenceNumber::from_u64(4))],
1084                        None
1085                    )
1086                ),
1087                (
1088                    certs[3].key(),
1089                    AssignedVersions::new_for_testing(
1090                        vec![((id, init_shared_version), SequenceNumber::from_u64(10))],
1091                        None
1092                    )
1093                ),
1094            ]
1095        );
1096    }
1097
1098    /// Generate a transaction that uses shared objects as specified in the parameters.
1099    /// Also uses a gas object with specified version.
1100    /// The version of the gas object is used to manipulate the lamport version of this transaction.
1101    fn generate_shared_objs_tx_with_gas_version(
1102        shared_objects: &[(ObjectID, SequenceNumber, bool)],
1103        gas_object_version: u64,
1104    ) -> VerifiedExecutableTransaction {
1105        let mut tx_builder = TestTransactionBuilder::new(
1106            SuiAddress::ZERO,
1107            (
1108                ObjectID::random(),
1109                SequenceNumber::from_u64(gas_object_version),
1110                ObjectDigest::random(),
1111            ),
1112            0,
1113        );
1114        let tx_data = {
1115            let builder = tx_builder.ptb_builder_mut();
1116            for (shared_object_id, shared_object_init_version, shared_object_mutable) in
1117                shared_objects
1118            {
1119                builder
1120                    .obj(ObjectArg::SharedObject {
1121                        id: *shared_object_id,
1122                        initial_shared_version: *shared_object_init_version,
1123                        mutability: if *shared_object_mutable {
1124                            SharedObjectMutability::Mutable
1125                        } else {
1126                            SharedObjectMutability::Immutable
1127                        },
1128                    })
1129                    .unwrap();
1130            }
1131            tx_builder.build()
1132        };
1133        let tx = SenderSignedData::new(tx_data, vec![]);
1134        VerifiedExecutableTransaction::new_unchecked(ExecutableTransaction::new_from_data_and_sig(
1135            tx,
1136            CertificateProof::new_system(0),
1137        ))
1138    }
1139
1140    struct WithdrawTestContext {
1141        authority: Arc<AuthorityState>,
1142        assignables: Vec<Schedulable<VerifiedExecutableTransaction>>,
1143        shared_objects: Vec<Object>,
1144    }
1145
1146    impl WithdrawTestContext {
1147        pub async fn new() -> Self {
1148            // Create a shared object for testing
1149            let shared_objects = vec![Object::shared_for_testing()];
1150            let mut config = ProtocolConfig::get_for_max_version_UNSAFE();
1151            config.set_enable_accumulators_for_testing(true);
1152            let authority = TestAuthorityBuilder::new()
1153                .with_starting_objects(&shared_objects)
1154                .with_protocol_config(config)
1155                .build()
1156                .await;
1157            Self {
1158                authority,
1159                assignables: vec![],
1160                shared_objects,
1161            }
1162        }
1163
1164        pub fn add_withdraw_transaction(&mut self) -> TransactionKey {
1165            // Generate random sender and gas object for each transaction
1166            let (sender, keypair) = get_account_key_pair();
1167            let gas_object = Object::with_owner_for_testing(sender);
1168            let gas_object_ref = gas_object.compute_object_reference();
1169            // Generate a unique gas price to make the transaction unique.
1170            let gas_price = (self.assignables.len() + 1) as u64;
1171            let mut tx_builder = TestTransactionBuilder::new(sender, gas_object_ref, gas_price);
1172            let tx_data = {
1173                let ptb_builder = tx_builder.ptb_builder_mut();
1174                ptb_builder
1175                    .funds_withdrawal(FundsWithdrawalArg::balance_from_sender(
1176                        200,
1177                        GAS::type_tag(),
1178                    ))
1179                    .unwrap();
1180                tx_builder.build()
1181            };
1182            let cert = VerifiedExecutableTransaction::new_for_testing(tx_data, &keypair);
1183            let key = cert.key();
1184            self.assignables.push(Schedulable::Transaction(cert));
1185            key
1186        }
1187
1188        pub fn add_settlement_transaction(&mut self) -> TransactionKey {
1189            let height = (self.assignables.len() + 1) as u64;
1190            let settlement = Schedulable::AccumulatorSettlement(0, height);
1191            let key = settlement.key();
1192            self.assignables.push(settlement);
1193            key
1194        }
1195
1196        pub fn add_withdraw_with_shared_object_transaction(&mut self) -> TransactionKey {
1197            // Generate random sender and gas object for each transaction
1198            let (sender, keypair) = get_account_key_pair();
1199            let gas_object = Object::with_owner_for_testing(sender);
1200            let gas_object_ref = gas_object.compute_object_reference();
1201            // Generate a unique gas price to make the transaction unique.
1202            let gas_price = (self.assignables.len() + 1) as u64;
1203            let mut tx_builder = TestTransactionBuilder::new(sender, gas_object_ref, gas_price);
1204            let tx_data = {
1205                let ptb_builder = tx_builder.ptb_builder_mut();
1206                // Add shared object to the transaction
1207                if let Some(shared_obj) = self.shared_objects.first() {
1208                    let id = shared_obj.id();
1209                    let init_version = shared_obj.owner.start_version().unwrap();
1210                    ptb_builder
1211                        .obj(ObjectArg::SharedObject {
1212                            id,
1213                            initial_shared_version: init_version,
1214                            mutability: SharedObjectMutability::Mutable,
1215                        })
1216                        .unwrap();
1217                }
1218                // Add balance withdraw
1219                ptb_builder
1220                    .funds_withdrawal(FundsWithdrawalArg::balance_from_sender(
1221                        200,
1222                        GAS::type_tag(),
1223                    ))
1224                    .unwrap();
1225                tx_builder.build()
1226            };
1227            let cert = VerifiedExecutableTransaction::new_for_testing(tx_data, &keypair);
1228            let key = cert.key();
1229            self.assignables.push(Schedulable::Transaction(cert));
1230            key
1231        }
1232
1233        pub fn assign_versions_from_consensus(&self) -> ConsensusSharedObjVerAssignment {
1234            let epoch_store = self.authority.epoch_store_for_testing();
1235            SharedObjVerManager::assign_versions_from_consensus(
1236                &epoch_store,
1237                self.authority.get_object_cache_reader().as_ref(),
1238                self.assignables.iter(),
1239                &BTreeMap::new(),
1240            )
1241            .unwrap()
1242        }
1243    }
1244
1245    #[tokio::test]
1246    async fn test_assign_versions_from_consensus_with_withdraws_simple() {
1247        // Note that we don't need a shared object to trigger withdraw version assignment.
1248        // In fact it is important that this works without a shared object.
1249        let mut ctx = WithdrawTestContext::new().await;
1250
1251        let acc_version = ctx
1252            .authority
1253            .get_object(&SUI_ACCUMULATOR_ROOT_OBJECT_ID)
1254            .unwrap()
1255            .version();
1256
1257        let withdraw_key = ctx.add_withdraw_transaction();
1258        let settlement_key = ctx.add_settlement_transaction();
1259
1260        let assigned_versions = ctx.assign_versions_from_consensus();
1261        assert_eq!(
1262            assigned_versions,
1263            ConsensusSharedObjVerAssignment {
1264                assigned_versions: AssignedTxAndVersions::new(vec![
1265                    (
1266                        withdraw_key,
1267                        AssignedVersions::new_for_testing(vec![], Some(acc_version))
1268                    ),
1269                    (
1270                        settlement_key,
1271                        AssignedVersions::new_for_testing(
1272                            vec![((SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version), acc_version)],
1273                            Some(acc_version)
1274                        )
1275                    ),
1276                ]),
1277                shared_input_next_versions: HashMap::from([(
1278                    (SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version),
1279                    acc_version.next()
1280                )]),
1281            }
1282        );
1283    }
1284
1285    #[tokio::test]
1286    async fn test_assign_versions_from_consensus_with_multiple_withdraws_and_settlements() {
1287        // Test with multiple withdrawals and multiple settlements, with settlement as the last transaction
1288        let mut ctx = WithdrawTestContext::new().await;
1289
1290        let acc_version = ctx
1291            .authority
1292            .get_object(&SUI_ACCUMULATOR_ROOT_OBJECT_ID)
1293            .unwrap()
1294            .version();
1295
1296        // First withdrawal and settlement
1297        let withdraw_key1 = ctx.add_withdraw_transaction();
1298        let settlement_key1 = ctx.add_settlement_transaction();
1299
1300        // Second withdrawal and settlement
1301        let withdraw_key2 = ctx.add_withdraw_transaction();
1302        let settlement_key2 = ctx.add_settlement_transaction();
1303
1304        // Third withdrawal and final settlement
1305        let withdraw_key3 = ctx.add_withdraw_transaction();
1306        let settlement_key3 = ctx.add_settlement_transaction();
1307
1308        let assigned_versions = ctx.assign_versions_from_consensus();
1309        assert_eq!(
1310            assigned_versions,
1311            ConsensusSharedObjVerAssignment {
1312                assigned_versions: AssignedTxAndVersions::new(vec![
1313                    (
1314                        withdraw_key1,
1315                        AssignedVersions::new_for_testing(vec![], Some(acc_version))
1316                    ),
1317                    (
1318                        settlement_key1,
1319                        AssignedVersions::new_for_testing(
1320                            vec![((SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version), acc_version)],
1321                            Some(acc_version)
1322                        )
1323                    ),
1324                    (
1325                        withdraw_key2,
1326                        AssignedVersions::new_for_testing(vec![], Some(acc_version.next()))
1327                    ),
1328                    (
1329                        settlement_key2,
1330                        AssignedVersions::new_for_testing(
1331                            vec![(
1332                                (SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version),
1333                                acc_version.next()
1334                            )],
1335                            Some(acc_version.next())
1336                        )
1337                    ),
1338                    (
1339                        withdraw_key3,
1340                        AssignedVersions::new_for_testing(vec![], Some(acc_version.next().next()))
1341                    ),
1342                    (
1343                        settlement_key3,
1344                        AssignedVersions::new_for_testing(
1345                            vec![(
1346                                (SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version),
1347                                acc_version.next().next()
1348                            )],
1349                            Some(acc_version.next().next())
1350                        )
1351                    ),
1352                ]),
1353                shared_input_next_versions: HashMap::from([(
1354                    (SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version),
1355                    acc_version.next().next().next()
1356                )]),
1357            }
1358        );
1359    }
1360
1361    #[tokio::test]
1362    async fn test_assign_versions_from_consensus_with_withdraw_and_shared_object() {
1363        // Test that a transaction can have both a withdrawal and use a shared object
1364        let mut ctx = WithdrawTestContext::new().await;
1365
1366        // Get the shared object info from the context
1367        let shared_obj_id = ctx.shared_objects[0].id();
1368        let shared_obj_version = ctx.shared_objects[0].owner.start_version().unwrap();
1369
1370        let acc_version = ctx
1371            .authority
1372            .get_object(&SUI_ACCUMULATOR_ROOT_OBJECT_ID)
1373            .unwrap()
1374            .version();
1375
1376        let withdraw_with_shared_key = ctx.add_withdraw_with_shared_object_transaction();
1377        let settlement_key = ctx.add_settlement_transaction();
1378
1379        let assigned_versions = ctx.assign_versions_from_consensus();
1380        assert_eq!(
1381            assigned_versions,
1382            ConsensusSharedObjVerAssignment {
1383                assigned_versions: AssignedTxAndVersions::new(vec![
1384                    (
1385                        withdraw_with_shared_key,
1386                        AssignedVersions::new_for_testing(
1387                            vec![((shared_obj_id, shared_obj_version), shared_obj_version)],
1388                            Some(acc_version)
1389                        )
1390                    ),
1391                    (
1392                        settlement_key,
1393                        AssignedVersions::new_for_testing(
1394                            vec![((SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version), acc_version)],
1395                            Some(acc_version)
1396                        )
1397                    ),
1398                ]),
1399                shared_input_next_versions: HashMap::from([
1400                    (
1401                        (SUI_ACCUMULATOR_ROOT_OBJECT_ID, acc_version),
1402                        acc_version.next()
1403                    ),
1404                    (
1405                        (shared_obj_id, shared_obj_version),
1406                        shared_obj_version.next()
1407                    ),
1408                ]),
1409            }
1410        );
1411    }
1412}