canic/memory/context/
subnet.rs

1use crate::{
2    Error,
3    cdk::structures::{DefaultMemoryImpl, cell::Cell, memory::VirtualMemory},
4    eager_static, ic_memory, impl_storable_bounded,
5    memory::{MemoryError, context::ContextError, id::context::SUBNET_CONTEXT_ID},
6};
7use candid::{CandidType, Principal};
8use serde::{Deserialize, Serialize};
9use std::cell::RefCell;
10use thiserror::Error as ThisError;
11
12//
13// SUBNET_CONTEXT
14//
15
16eager_static! {
17    static SUBNET_CONTEXT: RefCell<Cell<SubnetContextData, VirtualMemory<DefaultMemoryImpl>>> =
18        RefCell::new(Cell::init(
19            ic_memory!(SubnetState, SUBNET_CONTEXT_ID),
20            SubnetContextData::default(),
21        ));
22}
23
24///
25/// SubnetContextError
26///
27
28#[derive(Debug, ThisError)]
29pub enum SubnetContextError {
30    #[error("prime subnet pid has not been set")]
31    PrimeSubnetNotSet,
32
33    #[error("subnet pid has not been set")]
34    SubnetNotSet,
35}
36
37impl From<SubnetContextError> for Error {
38    fn from(err: SubnetContextError) -> Self {
39        MemoryError::from(ContextError::from(err)).into()
40    }
41}
42
43///
44/// SubnetContextData
45///
46
47#[derive(CandidType, Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
48pub struct SubnetContextData {
49    pub prime_subnet_pid: Option<Principal>,
50    pub subnet_pid: Option<Principal>,
51}
52
53impl_storable_bounded!(SubnetContextData, 64, true);
54
55///
56/// SubnetContext
57///
58
59pub struct SubnetContext;
60
61impl SubnetContext {
62    // ---- Subnet PID ----
63    #[must_use]
64    pub fn get_subnet_pid() -> Option<Principal> {
65        SUBNET_CONTEXT.with_borrow(|cell| cell.get().subnet_pid)
66    }
67
68    pub fn try_get_subnet_pid() -> Result<Principal, Error> {
69        Self::get_subnet_pid().ok_or_else(|| SubnetContextError::SubnetNotSet.into())
70    }
71
72    pub fn set_subnet_pid(pid: Principal) {
73        SUBNET_CONTEXT.with_borrow_mut(|cell| {
74            let mut data = *cell.get();
75            data.subnet_pid = Some(pid);
76            cell.set(data);
77        });
78    }
79
80    // ---- Prime Subnet PID ----
81    #[must_use]
82    pub fn get_prime_subnet_pid() -> Option<Principal> {
83        SUBNET_CONTEXT.with_borrow(|cell| cell.get().prime_subnet_pid)
84    }
85
86    pub fn try_get_prime_subnet_pid() -> Result<Principal, Error> {
87        Self::get_prime_subnet_pid().ok_or_else(|| SubnetContextError::PrimeSubnetNotSet.into())
88    }
89
90    pub fn set_prime_subnet_pid(pid: Principal) {
91        SUBNET_CONTEXT.with_borrow_mut(|cell| {
92            let mut data = *cell.get();
93            data.prime_subnet_pid = Some(pid);
94            cell.set(data);
95        });
96    }
97
98    // ---- Import / Export ----
99
100    #[must_use]
101    pub fn export() -> SubnetContextData {
102        SUBNET_CONTEXT.with_borrow(|cell| *cell.get())
103    }
104}
canic/memory/context/subnet.rs

canic/memory/context/
subnet.rs
