ab_client_consensus_common/
state.rs

1#[cfg(test)]
2mod tests;
3
4use ab_aligned_buffer::SharedAlignedBuffer;
5use ab_client_api::ContractSlotState;
6use ab_core_primitives::address::Address;
7use ab_core_primitives::hashes::Blake3Hash;
8use ab_core_primitives::trees::{StateOwnerContractsSparseMerkleTree, StateOwnersSparseMerkleTree};
9use ab_merkle_tree::sparse::Leaf;
10use blake3::hash;
11use std::collections::BTreeMap;
12use std::mem::MaybeUninit;
13use std::num::NonZeroU128;
14use triomphe::{Arc, UniqueArc};
15
16// TODO: This is very inefficient, but will do for now
17#[derive(Debug, Clone)]
18pub struct GlobalState {
19    /// Map from the owner address to a map from the management contract to the corresponding state
20    state: BTreeMap<Address, BTreeMap<Address, SharedAlignedBuffer>>,
21    total_len: usize,
22}
23
24impl GlobalState {
25    pub fn new(system_contract_states: &[ContractSlotState]) -> Self {
26        let mut state = BTreeMap::<Address, BTreeMap<_, _>>::new();
27        for system_contract_state in system_contract_states {
28            state
29                .entry(system_contract_state.owner)
30                .or_default()
31                .insert(
32                    system_contract_state.contract,
33                    system_contract_state.contents.clone(),
34                );
35        }
36
37        Self {
38            state,
39            total_len: system_contract_states.len(),
40        }
41    }
42
43    pub fn to_system_contract_states(&self) -> Arc<[ContractSlotState]> {
44        let mut system_contract_states =
45            UniqueArc::<[MaybeUninit<ContractSlotState>]>::new_uninit_slice(self.total_len);
46
47        self.state
48            .iter()
49            .flat_map(|(owner, state)| {
50                state.iter().map(|(contract, contents)| ContractSlotState {
51                    owner: *owner,
52                    contract: *contract,
53                    contents: contents.clone(),
54                })
55            })
56            .zip(system_contract_states.iter_mut())
57            .for_each(|(input, output)| {
58                output.write(input);
59            });
60
61        // SAFETY: Just initialized all entries, internal invariant guarantees that `self.total_len`
62        // matches the number of entries deep in the state tree
63        unsafe { UniqueArc::assume_init_slice(system_contract_states) }.shareable()
64    }
65
66    pub fn root(&self) -> Blake3Hash {
67        let mut previous_owner = None;
68
69        let maybe_state_root = StateOwnersSparseMerkleTree::compute_root_only(
70            self.state.iter().flat_map(|(&owner, state)| {
71                let owner = u128::from(owner);
72                // Owners are strictly increasing due to `BTreeMap`, hence no underflow
73                let skip_count = match previous_owner.replace(owner) {
74                    Some(previous_owner) => owner - previous_owner - 1,
75                    None => owner,
76                };
77                let skip_leaf =
78                    NonZeroU128::new(skip_count).map(|skip_count| Leaf::Empty { skip_count });
79
80                let mut previous_contract = None;
81
82                let maybe_owner_root = StateOwnerContractsSparseMerkleTree::compute_root_only(
83                    state.iter().flat_map(|(&contract, contents)| {
84                        let contract = u128::from(contract);
85                        // Contracts are strictly increasing due to `BTreeMap`, hence no underflow
86                        let skip_count = match previous_contract.replace(contract) {
87                            Some(previous_contract) => contract - previous_contract - 1,
88                            None => contract,
89                        };
90                        let skip_leaf = NonZeroU128::new(skip_count)
91                            .map(|skip_count| Leaf::Empty { skip_count });
92
93                        skip_leaf.into_iter().chain([Leaf::OccupiedOwned {
94                            leaf: *hash(contents.as_slice()).as_bytes(),
95                        }])
96                    }),
97                );
98                let owner_root = maybe_owner_root.expect(
99                    "The number of leaves is limited by address space, which is 128-bit; qed",
100                );
101
102                skip_leaf
103                    .into_iter()
104                    .chain([Leaf::OccupiedOwned { leaf: owner_root }])
105            }),
106        );
107        let state_root = maybe_state_root
108            .expect("The number of leaves is limited by address space, which is 128-bit; qed");
109
110        Blake3Hash::new(state_root)
111    }
112}
ab_client_consensus_common/state.rs

ab_client_consensus_common/
state.rs
