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.clone().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_insert_with(|| NodeSubstates::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                .iter()
173                .map(|(key, _value)| key.clone())
174                .take(count.try_into().unwrap())
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                .iter()
199                .map(|(key, _)| key.clone())
200                .take(count.try_into().unwrap())
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)| {
243                (
244                    k.clone(),
245                    v.iter().map(|(k, v)| (k.clone(), v.len())).collect(),
246                )
247            })
248            .collect();
249
250        self.nodes.insert(node_id, substates);
251
252        for (partition_number, partition) in sizes {
253            for (substate_key, substate_size) in partition {
254                on_io_access(
255                    self,
256                    IOAccess::HeapSubstateUpdated {
257                        canonical_substate_key: CanonicalSubstateKey {
258                            node_id,
259                            partition_number,
260                            substate_key: substate_key.clone(),
261                        },
262                        old_size: None,
263                        new_size: Some(substate_size),
264                    },
265                )?;
266            }
267        }
268
269        Ok(())
270    }
271
272    /// Removes node.
273    pub fn remove_node<E, F: FnMut(&Heap, IOAccess) -> Result<(), E>>(
274        &mut self,
275        node_id: &NodeId,
276        on_io_access: &mut F,
277    ) -> Result<NodeSubstates, CallbackError<HeapRemoveNodeError, E>> {
278        let node_substates = match self.nodes.remove(node_id) {
279            Some(node_substates) => node_substates,
280            None => Err(CallbackError::Error(HeapRemoveNodeError::NodeNotFound(
281                node_id.clone().into(),
282            )))?,
283        };
284
285        for (partition_number, partition) in &node_substates {
286            for (substate_key, substate_value) in partition {
287                on_io_access(
288                    self,
289                    IOAccess::HeapSubstateUpdated {
290                        canonical_substate_key: CanonicalSubstateKey {
291                            node_id: *node_id,
292                            partition_number: *partition_number,
293                            substate_key: substate_key.clone(),
294                        },
295                        old_size: Some(substate_value.len()),
296                        new_size: None,
297                    },
298                )
299                .map_err(CallbackError::CallbackError)?;
300            }
301        }
302
303        Ok(node_substates)
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn test_heap_size_accounting() {
313        let mut heap = Heap::new();
314        let mut total_size = 0;
315
316        let mut on_io_access = |_: &_, io_access| {
317            match io_access {
318                IOAccess::HeapSubstateUpdated {
319                    canonical_substate_key,
320                    old_size,
321                    new_size,
322                } => {
323                    if old_size.is_none() {
324                        total_size += canonical_substate_key.len();
325                    }
326                    if new_size.is_none() {
327                        total_size -= canonical_substate_key.len();
328                    }
329
330                    total_size += new_size.unwrap_or_default();
331                    total_size -= old_size.unwrap_or_default();
332                }
333                _ => {}
334            }
335
336            Result::<(), ()>::Ok(())
337        };
338
339        let node_id = NodeId([0u8; NodeId::LENGTH]);
340        let partition_number1 = PartitionNumber(5);
341        let partition_number2 = PartitionNumber(6);
342        let key1 = SubstateKey::Map(scrypto_encode(&"1").unwrap());
343        let key2 = SubstateKey::Map(scrypto_encode(&"2").unwrap());
344        heap.create_node(
345            NodeId([0u8; NodeId::LENGTH]),
346            btreemap!(
347                partition_number1 => btreemap!(
348                    key1.clone() => IndexedScryptoValue::from_typed("A"),
349                    key2.clone() => IndexedScryptoValue::from_typed("B"),
350                ),
351                partition_number2 => btreemap!(
352                    key1.clone() => IndexedScryptoValue::from_typed("C"),
353                    key2.clone() => IndexedScryptoValue::from_typed("D"),
354                )
355            ),
356            &mut on_io_access,
357        )
358        .unwrap();
359        heap.set_substate(
360            node_id,
361            partition_number1,
362            key1,
363            IndexedScryptoValue::from_typed("E"),
364            &mut on_io_access,
365        )
366        .unwrap();
367        heap.drain_substates(&node_id, partition_number1, 1, &mut on_io_access)
368            .unwrap();
369        heap.remove_substate(&node_id, partition_number2, &key2, &mut on_io_access)
370            .unwrap();
371        heap.remove_partition(&node_id, partition_number2, &mut on_io_access)
372            .unwrap();
373        heap.remove_node(&node_id, &mut on_io_access).unwrap();
374        assert_eq!(total_size, 0);
375    }
376}
radix_engine/kernel/heap.rs

radix_engine/kernel/
heap.rs
