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::num::NonZeroU128;
9use std::sync::Arc as StdArc;
10
11type Smt128 = SparseMerkleTree<{ size_of::<Address>() as u8 * u8::BITS as u8 }>;
12
13// TODO: This is very inefficient, but will do for now
14#[derive(Debug, Clone)]
15pub struct GlobalState {
16    /// Map from the owner address to a map from the management contract to the corresponding state
17    state: BTreeMap<Address, BTreeMap<Address, SharedAlignedBuffer>>,
18    total_len: usize,
19}
20
21impl GlobalState {
22    pub fn new(system_contract_states: &[ContractSlotState]) -> Self {
23        let mut state = BTreeMap::<Address, BTreeMap<_, _>>::new();
24        for system_contract_state in system_contract_states.iter() {
25            state
26                .entry(system_contract_state.owner)
27                .or_default()
28                .insert(
29                    system_contract_state.contract,
30                    system_contract_state.contents.clone(),
31                );
32        }
33
34        Self {
35            state,
36            total_len: system_contract_states.len(),
37        }
38    }
39
40    pub fn to_system_contract_states(&self) -> StdArc<[ContractSlotState]> {
41        let mut system_contract_states =
42            StdArc::<[ContractSlotState]>::new_uninit_slice(self.total_len);
43
44        self.state
45            .iter()
46            .flat_map(|(owner, state)| {
47                state.iter().map(|(contract, contents)| ContractSlotState {
48                    owner: *owner,
49                    contract: *contract,
50                    contents: contents.clone(),
51                })
52            })
53            // SAFETY: A single pointer and a single use
54            .zip(unsafe { StdArc::get_mut_unchecked(&mut system_contract_states) })
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 { system_contract_states.assume_init() }
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
