ab_system_contract_state/
lib.rs

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