ndata/
heap.rs

1#![cfg_attr(feature = "no_std_support", no_std)]
2
3extern crate alloc;
4use crate::usizemap::UsizeMap;
5
6// Use alloc::vec::Vec when the no_std_support feature is enabled
7#[cfg(feature="no_std_support")]
8use alloc::vec::Vec;
9// If std is available (default), Vec is used from the standard library prelude
10#[cfg(not(feature="no_std_support"))]
11use std::vec::Vec;
12
13use core::fmt::{self, Debug};
14
15// Internal struct to hold the data and its reference count.
16// Not public, as it's an implementation detail of Heap.
17#[derive(Debug)]
18struct Blob<T> {
19    data: T,
20    count: usize,
21}
22
23/// A reference counting container for objects of a given type with basic
24/// garbage collection based on reference counts reaching zero.
25///
26/// Keys are `usize` indices returned by `push`.
27pub struct Heap<T> {
28    data: UsizeMap<Blob<T>>,
29}
30
31// Implement Debug manually to avoid requiring T: Debug for the struct definition,
32// but require it for the Debug implementation.
33impl<T: Debug> Debug for Heap<T> {
34    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
35        f.debug_struct("Heap")
36         .field("data", &self.data)
37         .finish()
38    }
39}
40
41impl<T> Heap<T> {
42    /// Creates a new, empty `Heap`.
43    #[inline]
44    pub fn new() -> Heap<T> {
45        Heap {
46            data: UsizeMap::new(),
47        }
48    }
49
50    /// Pushes a value onto the heap, returning a stable `usize` key.
51    /// The initial reference count is 1.
52    pub fn push(&mut self, data: T) -> usize {
53        let blob = Blob {
54            data,
55            count: 1,
56        };
57        self.data.insert(blob)
58    }
59
60    /// Returns a mutable reference to the value associated with the given key.
61    ///
62    /// # Panics
63    /// Panics if the `index` is not a valid key.
64    pub fn get(&mut self, index: usize) -> &mut T {
65        &mut self.data.get_mut(index).expect("Heap::get: Invalid index").data
66    }
67
68    /// Returns a mutable reference to the value associated with the key, if it exists.
69    pub fn try_get(&mut self, index: usize) -> Option<&mut T> {
70        self.data.get_mut(index).map(|blob| &mut blob.data)
71    }
72
73    /// Checks if the heap contains a value for the specified key.
74    pub fn contains_key(&self, index: usize) -> bool {
75        self.data.contains_key(index)
76    }
77
78    /// Returns the current reference count for the value associated with the key.
79    /// Returns 0 if the key does not exist.
80    pub fn count(&mut self, index: usize) -> usize {
81        self.data.get_mut(index).map(|b| b.count).unwrap_or(0)
82    }
83
84    /// Increments the reference count for the value associated with the key.
85    ///
86    /// # Panics
87    /// Panics if the `index` is not valid.
88    pub fn incr(&mut self, index: usize) {
89        if let Some(blob) = self.data.get_mut(index) {
90            blob.count += 1;
91        } else {
92            panic!("Heap::incr: Invalid index {}", index);
93        }
94    }
95
96    /// Decrements the reference count. If it reaches zero, the value is removed.
97    ///
98    /// # Panics
99    /// Panics if the `index` is not valid.
100    pub fn decr(&mut self, index: usize) {
101        // Step 1: Check the count and decide if we need to remove.
102        // We CANNOT call remove() inside this block because self.data is borrowed.
103        let should_remove = if let Some(blob) = self.data.get_mut(index) {
104            if blob.count > 1 {
105                blob.count -= 1;
106                false
107            } else {
108                // count is 1 (or 0, which shouldn't happen), so we remove it.
109                true
110            }
111        } else {
112            panic!("Heap::decr: Invalid index {}", index);
113        };
114
115        // Step 2: Perform removal if necessary (mutable borrow has ended)
116        if should_remove {
117            self.data.remove(index);
118        }
119    }
120
121    /// Forcibly removes an item from the heap regardless of its reference count.
122    /// Returns the data if it existed.
123    ///
124    /// Useful for error recovery or session cleanup.
125    pub fn free(&mut self, index: usize) -> Option<T> {
126        self.data.remove(index).map(|b| b.data)
127    }
128
129    /// Returns a vector containing all the keys currently present.
130    pub fn keys(&self) -> Vec<usize> {
131        self.data.keys()
132    }
133
134    /// Clears the heap, removing all values.
135    pub fn clear(&mut self) {
136        self.data.clear();
137    }
138
139    /// Shrinks the underlying memory storage to fit the current data.
140    pub fn shrink_to_fit(&mut self) {
141        self.data.shrink_to_fit();
142    }
143}
144
145// Implement Default trait for Heap<T>
146impl<T> Default for Heap<T> {
147    fn default() -> Self {
148        Self::new()
149    }
150}
ndata/heap.rs

ndata/
heap.rs
