ab_client_consensus_common/
state.rs

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

ab_client_consensus_common/
state.rs
