ab_system_contract_address_allocator/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
#![no_std]

use ab_contracts_common::env::Env;
use ab_contracts_common::{Address, ContractError, ShardIndex};
use ab_contracts_io_type::trivial_type::TrivialType;
use ab_contracts_macros::contract;

#[derive(Copy, Clone, TrivialType)]
#[repr(C)]
pub struct AddressAllocator {
    /// Next address to be allocated on this shard
    pub next_address: u64,
    /// Max address to be allocated on this shard
    pub max_address: u64,
}

#[contract]
impl AddressAllocator {
    /// Initialize address allocator for a shard
    #[init]
    pub fn new(#[env] env: &Env) -> Self {
        let shard_index = env.shard_index();

        let expected_self_address = shard_index.to_u32() as u64 * ShardIndex::MAX_SHARDS as u64;
        debug_assert_eq!(
            env.own_address(),
            Address::from(expected_self_address),
            "Unexpected allocator address"
        );

        Self {
            next_address: expected_self_address + 1,
            max_address: (shard_index.to_u32() as u64 + 1) * ShardIndex::MAX_SHARDS as u64 - 1,
        }
    }

    /// Allocate a new address for a contract.
    ///
    /// This can only be called by [`Address::SYSTEM_CODE`] contract.
    #[update]
    pub fn allocate_address(&mut self, #[env] env: &mut Env) -> Result<Address, ContractError> {
        if env.caller() != Address::SYSTEM_CODE {
            return Err(ContractError::AccessDenied);
        }

        let next_address = self.next_address;
        if next_address == self.max_address {
            // No more addresses can be allocated on this shard
            return Err(ContractError::InvalidState);
        }

        self.next_address += 1;
        Ok(Address::from(next_address))
    }
}