ab_system_contract_address_allocator/
lib.rs

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

ab_system_contract_address_allocator/
lib.rs
