Skip to main content

ab_system_contract_state/
lib.rs

1#![no_std]
2
3use ab_contracts_common::ContractError;
4use ab_contracts_common::env::Env;
5use ab_contracts_macros::contract;
6use ab_core_primitives::address::Address;
7use ab_io_type::trivial_type::TrivialType;
8use ab_io_type::variable_bytes::VariableBytes;
9use core::mem::MaybeUninit;
10
11// TODO: How/where should this limit be defined?
12pub const RECOMMENDED_STATE_CAPACITY: u32 = 1024;
13
14/// Helper function that calls provided function with new empty state buffer
15#[inline]
16pub fn with_state_buffer<F, R>(f: F) -> R
17where
18    F: FnOnce(&mut VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>) -> R,
19{
20    let mut state_bytes = [MaybeUninit::uninit(); RECOMMENDED_STATE_CAPACITY as usize];
21    let mut state_size = 0;
22    let mut new_state = VariableBytes::from_uninit(&mut state_bytes, &mut state_size)
23        .expect("Size is zero and the buffer is smaller than 4 GiB; qed");
24    f(&mut new_state)
25}
26
27/// Helper function that calls provided function with new empty state buffer pair
28#[inline]
29pub fn with_state_buffer_pair<F, R>(f: F) -> R
30where
31    F: FnOnce(
32        &mut VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
33        &mut VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
34    ) -> R,
35{
36    let mut state_bytes_a = [MaybeUninit::uninit(); RECOMMENDED_STATE_CAPACITY as usize];
37    let mut state_size_a = 0;
38    let mut new_state_a = VariableBytes::from_uninit(&mut state_bytes_a, &mut state_size_a)
39        .expect("Size is zero and the buffer is smaller than 4 GiB; qed");
40
41    let mut state_bytes_b = [MaybeUninit::uninit(); RECOMMENDED_STATE_CAPACITY as usize];
42    let mut state_size_b = 0;
43    let mut new_state_b = VariableBytes::from_uninit(&mut state_bytes_b, &mut state_size_b)
44        .expect("Size is zero and the buffer is smaller than 4 GiB; qed");
45
46    f(&mut new_state_a, &mut new_state_b)
47}
48
49#[derive(Debug, Copy, Clone, TrivialType)]
50#[repr(C)]
51pub struct State;
52
53#[contract]
54impl State {
55    /// Initialize state.
56    ///
57    /// Similar to [`State::write()`], but returns error if the state is not empty.
58    #[update]
59    pub fn initialize(
60        #[env] env: &mut Env<'_>,
61        #[slot] (address, contract_state): (
62            &Address,
63            &mut VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
64        ),
65        #[input] state: &VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
66    ) -> Result<(), ContractError> {
67        if !Self::is_empty(contract_state) {
68            return Err(ContractError::Conflict);
69        }
70
71        Self::write(env, (address, contract_state), state)
72    }
73
74    /// Write state.
75    ///
76    /// Only direct caller is allowed to write its own state for security reasons.
77    #[update]
78    pub fn write(
79        #[env] env: &mut Env<'_>,
80        #[slot] (address, state): (&Address, &mut VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>),
81        #[input] new_state: &VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
82    ) -> Result<(), ContractError> {
83        // TODO: Check shard?
84        if env.caller() != address {
85            return Err(ContractError::Forbidden);
86        }
87
88        if !state.copy_from(new_state) {
89            return Err(ContractError::BadInput);
90        }
91
92        Ok(())
93    }
94
95    /// Compare state with a given old state and write new state if old state matches.
96    ///
97    /// Only direct caller is allowed to write its own state for security reasons.
98    ///
99    /// Returns boolean indicating whether write happened or not.
100    #[update]
101    pub fn compare_and_write(
102        #[env] env: &mut Env<'_>,
103        #[slot] (address, state): (&Address, &mut VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>),
104        #[input] old_state: &VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
105        #[input] new_state: &VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
106    ) -> Result<bool, ContractError> {
107        // TODO: Check shard?
108        if env.caller() != address {
109            return Err(ContractError::Forbidden);
110        }
111
112        if state.get_initialized() != old_state.get_initialized() {
113            return Ok(false);
114        }
115
116        if !state.copy_from(new_state) {
117            return Err(ContractError::BadInput);
118        }
119
120        Ok(true)
121    }
122
123    /// Read state
124    #[view]
125    pub fn read(
126        #[slot] contract_state: &VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
127        #[output] state: &mut VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
128    ) -> Result<(), ContractError> {
129        if state.copy_from(contract_state) {
130            Ok(())
131        } else {
132            Err(ContractError::BadInput)
133        }
134    }
135
136    /// Check if the state is empty
137    #[view]
138    pub fn is_empty(
139        #[slot] contract_state: &VariableBytes<'_, RECOMMENDED_STATE_CAPACITY>,
140    ) -> bool {
141        contract_state.size() == 0
142    }
143}