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 =
27            u128::from(shard_index.as_u32()) * ShardIndex::MAX_ADDRESSES_PER_SHARD.get();
28        debug_assert_eq!(
29            env.own_address(),
30            Address::from(expected_self_address),
31            "Unexpected allocator address"
32        );
33
34        Self {
35            next_address: expected_self_address + 1,
36            max_address: expected_self_address + ShardIndex::MAX_ADDRESSES_PER_SHARD.get() - 1,
37        }
38    }
39
40    /// Allocate a new address for a contract.
41    ///
42    /// This can only be called by [`Address::SYSTEM_CODE`] contract.
43    #[update]
44    pub fn allocate_address(&mut self, #[env] env: &mut Env<'_>) -> Result<Address, ContractError> {
45        if env.caller() != Address::SYSTEM_CODE {
46            return Err(ContractError::Forbidden);
47        }
48
49        let next_address = self.next_address;
50        if next_address >= self.max_address {
51            // No more addresses can be allocated on this shard
52            return Err(ContractError::Forbidden);
53        }
54
55        self.next_address += 1;
56        Ok(Address::from(next_address))
57    }
58}
ab_system_contract_address_allocator/lib.rs

ab_system_contract_address_allocator/
lib.rs
