ic_canister_kit/
stable.rs

1use std::cell::RefCell;
2
3use ic_stable_structures::{DefaultMemoryImpl, memory_manager::MemoryManager};
4
5/// 简化虚拟内存
6pub type VirtualMemory = ic_stable_structures::memory_manager::VirtualMemory<ic_stable_structures::DefaultMemoryImpl>;
7pub use ic_stable_structures::GrowFailed;
8pub use ic_stable_structures::Memory;
9pub use ic_stable_structures::Storable;
10pub use ic_stable_structures::memory_manager::MemoryId;
11pub use ic_stable_structures::storable::Bound;
12pub use ic_stable_structures::writer::Writer;
13pub use std::borrow::Cow;
14
15/// 稳定对象
16/// ! 读取和写入都是全量操作，成本比较大
17pub type StableCell<T> = ic_stable_structures::Cell<T, VirtualMemory>;
18/// 稳定列表
19/// ! 存储有限长度数据，若不固定长度，则按照最大长度存储，不均匀的数据使用空间浪费比较严重
20/// ! push 和 pop 没有任意位置删除的功能
21/// ! 若不在乎顺序，则移动末尾元素到被删除的位置可实现任意删除。结合 StableBTreeMap 存储双向的索引数据，可实现任意位置删除。
22/// ! 最大的问题还是数据长度问题，任意删除功能不是核心难题。
23pub type StableVec<T> = ic_stable_structures::Vec<T, VirtualMemory>;
24/// 稳定映射
25pub type StableBTreeMap<K, V> = ic_stable_structures::BTreeMap<K, V, VirtualMemory>;
26/// 稳定日志
27/// ! 支持变长 无法删除和移动
28pub type StableLog<T> = ic_stable_structures::Log<T, VirtualMemory, VirtualMemory>;
29/// 稳定优先级队列 按照排序方式存放数据
30/// ! 内部使用 vec 方式实现，优缺点一致
31pub type StablePriorityQueue<T> = ic_stable_structures::MinHeap<T, VirtualMemory>;
32
33thread_local! {
34    static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> = RefCell::new(MemoryManager::init(DefaultMemoryImpl::default()));
35}
36
37// 最大支持 255 个内存片段
38// 这里将第 254 号用作保存堆内存序列化存档数据的内存，业务数据等不应该使用 254 号
39// https://github.com/dfinity/stable-structures/blob/29fad0d7b86333527b41924582bfede5bdc5dfc1/src/memory_manager.rs#L55
40const MEMORY_ID_UPGRADED: MemoryId = MemoryId::new(254);
41
42/// 获取虚拟内存
43#[inline]
44pub fn get_virtual_memory(memory_id: MemoryId) -> VirtualMemory {
45    MEMORY_MANAGER.with(|memory_manager| memory_manager.borrow().get(memory_id))
46}
47
48/// 获取升级用的虚拟内存
49#[inline]
50pub fn get_upgrades_memory() -> VirtualMemory {
51    get_virtual_memory(MEMORY_ID_UPGRADED)
52}
53
54/// 初始化内存
55pub fn init_cell_data<T: Storable>(memory_id: MemoryId, default: T) -> StableCell<T> {
56    StableCell::init(get_virtual_memory(memory_id), default)
57}
58/// 初始化内存
59pub fn init_vec_data<T: Storable>(memory_id: MemoryId) -> StableVec<T> {
60    StableVec::init(get_virtual_memory(memory_id))
61}
62/// 初始化内存
63pub fn init_map_data<K: Storable + Ord + Clone, V: Storable>(memory_id: MemoryId) -> StableBTreeMap<K, V> {
64    StableBTreeMap::init(get_virtual_memory(memory_id))
65}
66/// 初始化内存
67pub fn init_log_data<T: Storable>(index_memory_id: MemoryId, data_memory_id: MemoryId) -> StableLog<T> {
68    StableLog::init(get_virtual_memory(index_memory_id), get_virtual_memory(data_memory_id))
69}
70/// 初始化内存
71pub fn init_priority_queue_data<T: Storable + PartialOrd>(memory_id: MemoryId) -> StablePriorityQueue<T> {
72    StablePriorityQueue::init(get_virtual_memory(memory_id))
73}
74
75/// 包装升级内存
76pub struct WriteUpgradeMemory<'a, M> {
77    writer: Writer<'a, M>,
78}
79
80/// 包装升级内存
81pub struct ReadUpgradeMemory<'a, M> {
82    memory: &'a M,
83    offset: u64,
84}
85
86impl<'a, M: Memory> WriteUpgradeMemory<'a, M> {
87    /// 构造升级对象
88    pub fn new(memory: &'a mut M) -> WriteUpgradeMemory<'a, M> {
89        Self {
90            writer: Writer::new(memory, 0),
91        }
92    }
93
94    /// 写入升级数据
95    pub fn write(&mut self, bytes: &[u8]) -> Result<(), GrowFailed> {
96        self.writer.write(bytes)
97    }
98
99    /// 写入 u32
100    pub fn write_u32(&mut self, value: u32) -> Result<(), GrowFailed> {
101        let mut bytes = Vec::with_capacity(4);
102        common::u32_to_bytes(&mut bytes, value);
103        self.write(&bytes)
104    }
105
106    /// 写入 u64
107    pub fn write_u64(&mut self, value: u64) -> Result<(), GrowFailed> {
108        let mut bytes = Vec::with_capacity(8);
109        common::u64_to_bytes(&mut bytes, value);
110        self.write(&bytes)
111    }
112}
113
114impl<'a, M: Memory> ReadUpgradeMemory<'a, M> {
115    /// 构造升级对象
116    pub fn new(memory: &'a M) -> ReadUpgradeMemory<'a, M> {
117        Self { memory, offset: 0 }
118    }
119
120    /// 读取升级数据
121    pub fn read(&mut self, bytes: &mut [u8]) {
122        self.memory.read(self.offset, bytes);
123        self.offset += bytes.len() as u64;
124    }
125
126    /// 读取 u32
127    pub fn read_u32(&mut self) -> u32 {
128        let mut bytes = [0; 4];
129        self.read(&mut bytes);
130        common::u32_from_bytes(&bytes)
131    }
132
133    /// 读取 u64
134    pub fn read_u64(&mut self) -> u64 {
135        let mut bytes = [0; 8];
136        self.read(&mut bytes);
137        common::u64_from_bytes(&bytes)
138    }
139}
140
141/// 一些可能用到的工具方法
142pub mod common {
143    use super::*;
144
145    /// usize -> 4 bytes
146    #[inline]
147    pub fn usize_to_4bytes(buf: &mut Vec<u8>, value: usize) {
148        buf.extend(&(value as u32).to_bytes()[..]);
149    }
150
151    /// 4 bytes -> usize
152    #[inline]
153    pub fn usize_from_4bytes(bytes: &[u8]) -> usize {
154        u32::from_bytes(Cow::Borrowed(&bytes[..4])) as usize
155    }
156
157    /// u32 -> 4 bytes
158    #[inline]
159    pub fn u32_to_bytes(buf: &mut Vec<u8>, value: u32) {
160        buf.extend(&value.to_bytes()[..]);
161    }
162
163    /// 4 bytes -> u32
164    #[inline]
165    pub fn u32_from_bytes(bytes: &[u8]) -> u32 {
166        u32::from_bytes(Cow::Borrowed(&bytes[..4]))
167    }
168
169    /// u64 -> 8 bytes
170    #[inline]
171    pub fn u64_to_bytes(buf: &mut Vec<u8>, value: u64) {
172        buf.extend(&value.to_bytes()[..]);
173    }
174
175    /// 8 bytes -> u64
176    #[inline]
177    pub fn u64_from_bytes(bytes: &[u8]) -> u64 {
178        u64::from_bytes(Cow::Borrowed(&bytes[..8]))
179    }
180}
ic_canister_kit/stable.rs

ic_canister_kit/
stable.rs
