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