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