radix_engine/kernel/
heap.rs

1use crate::internal_prelude::*;
2use crate::track::interface::IOAccess;
3use crate::track::interface::{CallbackError, CanonicalSubstateKey, NodeSubstates};
4
5pub struct Heap {
6    nodes: NonIterMap<NodeId, NodeSubstates>,
7}
8
9#[derive(Debug, Clone, PartialEq, Eq, ScryptoSbor)]
10pub enum HeapRemovePartitionError {
11    NodeNotFound(error_models::ReferencedNodeId),
12    ModuleNotFound(PartitionNumber),
13}
14
15#[derive(Debug, Clone, PartialEq, Eq, ScryptoSbor)]
16pub enum HeapRemoveNodeError {
17    NodeNotFound(error_models::ReferencedNodeId),
18}
19
20impl Heap {
21    pub fn new() -> Self {
22        Self {
23            nodes: NonIterMap::new(),
24        }
25    }
26
27    pub fn is_empty(&self) -> bool {
28        self.nodes.is_empty()
29    }
30
31    pub fn remove_partition<E, F: FnMut(&Heap, IOAccess) -> Result<(), E>>(
32        &mut self,
33        node_id: &NodeId,
34        partition_number: PartitionNumber,
35        on_io_access: &mut F,
36    ) -> Result<
37        BTreeMap<SubstateKey, IndexedScryptoValue>,
38        CallbackError<HeapRemovePartitionError, E>,
39    > {
40        if let Some(substates) = self.nodes.get_mut(node_id) {
41            let partition = substates
42                .remove(&partition_number)
43                .ok_or(CallbackError::Error(
44                    HeapRemovePartitionError::ModuleNotFound(partition_number),
45                ))?;
46
47            for (substate_key, substate_value) in &partition {
48                on_io_access(
49                    self,
50                    IOAccess::HeapSubstateUpdated {
51                        canonical_substate_key: CanonicalSubstateKey {
52                            node_id: *node_id,
53                            partition_number,
54                            substate_key: substate_key.clone(),
55                        },
56                        old_size: Some(substate_value.len()),
57                        new_size: None,
58                    },
59                )
60                .map_err(CallbackError::CallbackError)?;
61            }
62
63            Ok(partition)
64        } else {
65            Err(CallbackError::Error(
66                HeapRemovePartitionError::NodeNotFound((*node_id).into()),
67            ))
68        }
69    }
70
71    /// Reads a substate
72    pub fn get_substate(
73        &self,
74        node_id: &NodeId,
75        partition_num: PartitionNumber,
76        substate_key: &SubstateKey,
77    ) -> Option<&IndexedScryptoValue> {
78        self.nodes
79            .get(node_id)
80            .and_then(|node| node.get(&partition_num))
81            .and_then(|partition_substates| partition_substates.get(substate_key))
82    }
83
84    /// Inserts or overwrites a substate
85    pub fn set_substate<E, F: FnMut(&Heap, IOAccess) -> Result<(), E>>(
86        &mut self,
87        node_id: NodeId,
88        partition_number: PartitionNumber,
89        substate_key: SubstateKey,
90        substate_value: IndexedScryptoValue,
91        on_io_access: &mut F,
92    ) -> Result<(), E> {
93        let entry = self
94            .nodes
95            .entry(node_id)
96            .or_default()
97            .entry(partition_number)
98            .or_default()
99            .entry(substate_key.clone());
100
101        let old_size;
102        let new_size = Some(substate_value.len());
103        match entry {
104            btree_map::Entry::Vacant(e) => {
105                old_size = None;
106                e.insert(substate_value);
107            }
108            btree_map::Entry::Occupied(mut e) => {
109                old_size = Some(e.get().len());
110                e.insert(substate_value);
111            }
112        }
113
114        on_io_access(
115            self,
116            IOAccess::HeapSubstateUpdated {
117                canonical_substate_key: CanonicalSubstateKey {
118                    node_id,
119                    partition_number,
120                    substate_key,
121                },
122                old_size,
123                new_size,
124            },
125        )?;
126
127        Ok(())
128    }
129
130    pub fn remove_substate<E, F: FnMut(&Heap, IOAccess) -> Result<(), E>>(
131        &mut self,
132        node_id: &NodeId,
133        partition_number: PartitionNumber,
134        substate_key: &SubstateKey,
135        on_io_access: &mut F,
136    ) -> Result<Option<IndexedScryptoValue>, E> {
137        let substate_value = self
138            .nodes
139            .get_mut(node_id)
140            .and_then(|n| n.get_mut(&partition_number))
141            .and_then(|s| s.remove(substate_key));
142
143        if let Some(value) = &substate_value {
144            on_io_access(
145                self,
146                IOAccess::HeapSubstateUpdated {
147                    canonical_substate_key: CanonicalSubstateKey {
148                        node_id: *node_id,
149                        partition_number,
150                        substate_key: substate_key.clone(),
151                    },
152                    old_size: Some(value.len()),
153                    new_size: None,
154                },
155            )?;
156        }
157
158        Ok(substate_value)
159    }
160
161    /// Scans the keys of a node's partition. On an non-existing node/partition, this
162    /// will return an empty vector
163    pub fn scan_keys(
164        &self,
165        node_id: &NodeId,
166        partition_num: PartitionNumber,
167        count: u32,
168    ) -> Vec<SubstateKey> {
169        let node_substates = self.nodes.get(node_id).and_then(|n| n.get(&partition_num));
170        if let Some(substates) = node_substates {
171            let substate_keys: Vec<SubstateKey> = substates
172                .keys()
173                .take(count.try_into().unwrap())
174                .cloned()
175                .collect();
176
177            substate_keys
178        } else {
179            vec![]
180        }
181    }
182
183    /// Drains the substates from a node's partition. On an non-existing node/partition, this
184    /// will return an empty vector
185    pub fn drain_substates<E, F: FnMut(&Heap, IOAccess) -> Result<(), E>>(
186        &mut self,
187        node_id: &NodeId,
188        partition_number: PartitionNumber,
189        count: u32,
190        on_io_access: &mut F,
191    ) -> Result<Vec<(SubstateKey, IndexedScryptoValue)>, E> {
192        let node_substates = self
193            .nodes
194            .get_mut(node_id)
195            .and_then(|n| n.get_mut(&partition_number));
196        if let Some(substates) = node_substates {
197            let keys: Vec<SubstateKey> = substates
198                .keys()
199                .take(count.try_into().unwrap())
200                .cloned()
201                .collect();
202
203            let mut items = Vec::new();
204
205            for key in keys {
206                let value = substates.remove(&key).unwrap();
207                items.push((key, value));
208            }
209
210            for (key, value) in &items {
211                on_io_access(
212                    self,
213                    IOAccess::HeapSubstateUpdated {
214                        canonical_substate_key: CanonicalSubstateKey {
215                            node_id: *node_id,
216                            partition_number,
217                            substate_key: key.clone(),
218                        },
219                        old_size: Some(value.len()),
220                        new_size: None,
221                    },
222                )?;
223            }
224
225            Ok(items)
226        } else {
227            Ok(vec![])
228        }
229    }
230
231    /// Inserts a new node to heap.
232    pub fn create_node<E, F: FnMut(&Heap, IOAccess) -> Result<(), E>>(
233        &mut self,
234        node_id: NodeId,
235        substates: NodeSubstates,
236        on_io_access: &mut F,
237    ) -> Result<(), E> {
238        assert!(!self.nodes.contains_key(&node_id));
239
240        let sizes: IndexMap<PartitionNumber, IndexMap<SubstateKey, usize>> = substates
241            .iter()
242            .map(|(k, v)| (*k, v.iter().map(|(k, v)| (k.clone(), v.len())).collect()))
243            .collect();
244
245        self.nodes.insert(node_id, substates);
246
247        for (partition_number, partition) in sizes {
248            for (substate_key, substate_size) in partition {
249                on_io_access(
250                    self,
251                    IOAccess::HeapSubstateUpdated {
252                        canonical_substate_key: CanonicalSubstateKey {
253                            node_id,
254                            partition_number,
255                            substate_key: substate_key.clone(),
256                        },
257                        old_size: None,
258                        new_size: Some(substate_size),
259                    },
260                )?;
261            }
262        }
263
264        Ok(())
265    }
266
267    /// Removes node.
268    pub fn remove_node<E, F: FnMut(&Heap, IOAccess) -> Result<(), E>>(
269        &mut self,
270        node_id: &NodeId,
271        on_io_access: &mut F,
272    ) -> Result<NodeSubstates, CallbackError<HeapRemoveNodeError, E>> {
273        let node_substates = match self.nodes.remove(node_id) {
274            Some(node_substates) => node_substates,
275            None => Err(CallbackError::Error(HeapRemoveNodeError::NodeNotFound(
276                (*node_id).into(),
277            )))?,
278        };
279
280        for (partition_number, partition) in &node_substates {
281            for (substate_key, substate_value) in partition {
282                on_io_access(
283                    self,
284                    IOAccess::HeapSubstateUpdated {
285                        canonical_substate_key: CanonicalSubstateKey {
286                            node_id: *node_id,
287                            partition_number: *partition_number,
288                            substate_key: substate_key.clone(),
289                        },
290                        old_size: Some(substate_value.len()),
291                        new_size: None,
292                    },
293                )
294                .map_err(CallbackError::CallbackError)?;
295            }
296        }
297
298        Ok(node_substates)
299    }
300}
301
302impl Default for Heap {
303    fn default() -> Self {
304        Self::new()
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    #[test]
313    fn test_heap_size_accounting() {
314        let mut heap = Heap::new();
315        let mut total_size = 0;
316
317        let mut on_io_access = |_: &_, io_access| {
318            if let IOAccess::HeapSubstateUpdated {
319                canonical_substate_key,
320                old_size,
321                new_size,
322            } = io_access
323            {
324                if old_size.is_none() {
325                    total_size += canonical_substate_key.len();
326                }
327                if new_size.is_none() {
328                    total_size -= canonical_substate_key.len();
329                }
330
331                total_size += new_size.unwrap_or_default();
332                total_size -= old_size.unwrap_or_default();
333            }
334
335            Result::<(), ()>::Ok(())
336        };
337
338        let node_id = NodeId([0u8; NodeId::LENGTH]);
339        let partition_number1 = PartitionNumber(5);
340        let partition_number2 = PartitionNumber(6);
341        let key1 = SubstateKey::Map(scrypto_encode(&"1").unwrap());
342        let key2 = SubstateKey::Map(scrypto_encode(&"2").unwrap());
343        heap.create_node(
344            NodeId([0u8; NodeId::LENGTH]),
345            btreemap!(
346                partition_number1 => btreemap!(
347                    key1.clone() => IndexedScryptoValue::from_typed("A"),
348                    key2.clone() => IndexedScryptoValue::from_typed("B"),
349                ),
350                partition_number2 => btreemap!(
351                    key1.clone() => IndexedScryptoValue::from_typed("C"),
352                    key2.clone() => IndexedScryptoValue::from_typed("D"),
353                )
354            ),
355            &mut on_io_access,
356        )
357        .unwrap();
358        heap.set_substate(
359            node_id,
360            partition_number1,
361            key1,
362            IndexedScryptoValue::from_typed("E"),
363            &mut on_io_access,
364        )
365        .unwrap();
366        heap.drain_substates(&node_id, partition_number1, 1, &mut on_io_access)
367            .unwrap();
368        heap.remove_substate(&node_id, partition_number2, &key2, &mut on_io_access)
369            .unwrap();
370        heap.remove_partition(&node_id, partition_number2, &mut on_io_access)
371            .unwrap();
372        heap.remove_node(&node_id, &mut on_io_access).unwrap();
373        assert_eq!(total_size, 0);
374    }
375}
radix_engine/kernel/heap.rs

radix_engine/kernel/
heap.rs
