ab_system_contract_address_allocator/
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_core_primitives::shard::ShardIndex;
8use ab_io_type::trivial_type::TrivialType;
9
10#[derive(Debug, Copy, Clone, TrivialType)]
11#[repr(C)]
12pub struct AddressAllocator {
13    /// Next address to be allocated on this shard
14    pub next_address: u128,
15    /// Max address to be allocated on this shard
16    pub max_address: u128,
17}
18
19#[contract]
20impl AddressAllocator {
21    /// Initialize address allocator for a shard
22    #[init]
23    pub fn new(#[env] env: &Env<'_>) -> Self {
24        let shard_index = env.shard_index();
25
26        let expected_self_address = u128::from(Address::system_address_allocator(shard_index));
27        debug_assert_eq!(
28            env.own_address(),
29            Address::from(expected_self_address),
30            "Unexpected allocator address"
31        );
32
33        Self {
34            next_address: expected_self_address + 1,
35            max_address: expected_self_address + (ShardIndex::MAX_ADDRESSES_PER_SHARD.get() - 1),
36        }
37    }
38
39    /// Allocate a new address for a contract.
40    ///
41    /// This can only be called by [`Address::SYSTEM_CODE`] contract.
42    #[update]
43    pub fn allocate_address(&mut self, #[env] env: &mut Env<'_>) -> Result<Address, ContractError> {
44        if env.caller() != Address::SYSTEM_CODE {
45            return Err(ContractError::Forbidden);
46        }
47
48        let next_address = self.next_address;
49        if next_address >= self.max_address {
50            // No more addresses can be allocated on this shard
51            return Err(ContractError::Forbidden);
52        }
53
54        self.next_address += 1;
55        Ok(Address::from(next_address))
56    }
57}
ab_system_contract_address_allocator/lib.rs

ab_system_contract_address_allocator/
lib.rs
