gc_lite/
heap.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright (c) 2025-2026 John Ray <996351336@qq.com>
3
4use std::{collections::HashMap, ptr::NonNull};
5
6use crate::{
7    gctype::GcTypeRegistry,
8    node::{GcHead, GcNodeFlag},
9    partition::{GcPartition, GcPartitionId},
10    scope::GcScopeState,
11};
12
13pub struct GcHeap {
14    /// Registered GC data type info
15    pub(super) node_dtypes: &'static GcTypeRegistry,
16
17    /// Partition management
18    pub(super) partitions: HashMap<GcPartitionId, GcPartition>,
19    /// Global memory usage limit for the entire heap, 0 for unlimited
20    pub(super) memory_limit: usize,
21    /// Global automatic GC threshold for the entire heap, 0 for disabled
22    pub(super) gc_threshold: usize,
23    /// Total memory used across all partitions
24    pub(super) total_memory_used: usize,
25    pub(crate) scope_stack: Vec<GcScopeState<'static>>, // DON'T use SmallVec here
26    /// Weak reference list, each slot stores (version, GcHeader)
27    pub(super) weak_slots: Vec<(u16, Option<NonNull<GcHead>>)>,
28
29    /// User provided opaque raw pointer
30    opaque: *mut u8,
31
32    #[cfg(debug_assertions)]
33    pub(crate) dbg_dropping_root_partition: Option<GcPartitionId>,
34    #[cfg(debug_assertions)]
35    pub(crate) dbg_living_nodes: std::collections::HashSet<NonNull<GcHead>>,
36}
37
38impl Drop for GcHeap {
39    fn drop(&mut self) {
40        // heap world is gone, dealloc all nodes live in it, regardless their status.
41        log::trace!("[heap::drop]");
42
43        for s in self.scope_stack.drain(..) {
44            unsafe {
45                s.abort();
46            }
47        }
48
49        let pars = std::mem::take(&mut self.partitions);
50        for (_, partition) in pars {
51            self.dispose_all_nodes(partition.nodes, Self::DUMMY_DISPOSE_CALLBACK);
52        }
53
54        #[cfg(debug_assertions)]
55        debug_assert!(
56            self.dbg_living_nodes.is_empty(),
57            "[O.o][heap drop] leaked nodes {:?}",
58            self.dbg_living_nodes
59        );
60    }
61}
62
63impl GcHeap {
64    pub const DUMMY_DISPOSE_CALLBACK: fn(&GcHeap, &GcHead) = |_, _| {};
65
66    /// Create a new garbage collection heap with an explicit GC type registry
67    pub fn new(registry: &'static GcTypeRegistry) -> Self {
68        Self {
69            partitions: HashMap::new(),
70            memory_limit: 0,
71            gc_threshold: 0,
72            total_memory_used: 0,
73            weak_slots: Vec::new(),
74            opaque: std::ptr::null_mut(),
75            node_dtypes: registry,
76            scope_stack: Vec::with_capacity(16),
77
78            #[cfg(debug_assertions)]
79            dbg_dropping_root_partition: None,
80            #[cfg(debug_assertions)]
81            dbg_living_nodes: std::collections::HashSet::with_capacity(128),
82        }
83    }
84
85    #[inline(always)]
86    pub const fn opaque(&self) -> *mut u8 {
87        self.opaque
88    }
89
90    pub const fn set_opaque(&mut self, opaque: *mut u8) {
91        self.opaque = opaque;
92    }
93
94    pub fn memory_limit(&self) -> usize {
95        self.memory_limit
96    }
97
98    pub fn set_memory_limit(&mut self, limit: usize) -> usize {
99        if limit == 0 {
100            self.memory_limit = 0;
101        } else {
102            let used = self.total_memory_used;
103            let applied = std::cmp::max(used, limit);
104            self.memory_limit = applied;
105
106            if self.gc_threshold > 0 && self.gc_threshold >= applied {
107                let adjusted = applied - (applied >> 2);
108                self.gc_threshold = adjusted;
109            }
110        }
111
112        self.memory_limit
113    }
114
115    pub fn gc_threshold(&self) -> usize {
116        self.gc_threshold
117    }
118
119    pub fn set_gc_threshold(&mut self, threshold: usize) -> usize {
120        if threshold > 0 && self.memory_limit > 0 {
121            let capped = self.memory_limit.saturating_mul(8).saturating_div(10);
122            self.gc_threshold = std::cmp::min(threshold, capped);
123        } else {
124            self.gc_threshold = threshold;
125        }
126
127        self.gc_threshold
128    }
129
130    /// Check if garbage collection is needed
131    #[inline(always)]
132    pub fn should_gc(&self) -> bool {
133        // If GC threshold > 0 and memory usage reaches threshold, trigger GC
134        // gc_threshold = 0 means automatic GC is disabled
135        self.gc_threshold > 0 && self.total_memory_used >= self.gc_threshold
136    }
137
138    /// Attach a node to partition's nodes chain.
139    ///
140    /// # Note
141    ///
142    /// This method **DO NOT** increase partitions' mem_use.
143    pub(crate) fn attach_node(&mut self, partition_id: GcPartitionId, mut node: NonNull<GcHead>) {
144        debug_assert!(!partition_id.is_null());
145
146        let n = unsafe { node.as_mut() };
147        debug_assert!(n.partition_id().is_null());
148        debug_assert!(n.next.is_none());
149        n.set_partition_id(partition_id);
150
151        let par = self.partitions.get_mut(&partition_id).unwrap();
152        par.nodes.prepend(node);
153    }
154
155    pub fn set_root_node(&mut self, mut node: NonNull<GcHead>) {
156        let n = unsafe { node.as_mut() };
157        if !n.is_root() {
158            let partition_id = n.partition_id();
159            if let Some(p) = self.partitions.get_mut(&partition_id) {
160                n.insert_flag(GcNodeFlag::ROOT);
161                if p.is_marking() {
162                    p.add_gray_node(node);
163                }
164            }
165        }
166    }
167
168    /// Check if `node` was allocated in this heap
169    pub fn contains(&self, node: NonNull<GcHead>) -> bool {
170        self.nodes(unsafe { node.as_ref().partition_id() })
171            .any(|p| p == node)
172    }
173
174    /// Protect node from being gc collected.
175    ///
176    /// 1. if node is local or root, it's protected, returns true
177    /// 1. otherwise if has current scope, add node to current scope and returns true
178    /// 1. can't protect, returns false
179    pub fn protect_node(&mut self, node: NonNull<GcHead>) -> bool {
180        self.current_scope().is_some_and(|s| s.add_non_local(node))
181    }
182
183    /// Protect nodes from being gc collected, for each node do following steps:
184    ///
185    /// 1. if node is local or root, do nothing
186    /// 1. if has current scope, add node to current scope
187    /// 1. can't protect, returns false
188    pub fn protect_nodes_iter(&mut self, nodes: impl Iterator<Item = NonNull<GcHead>>) {
189        if let Some(s) = self.current_scope() {
190            for n in nodes {
191                s.add_non_local(n);
192            }
193        }
194    }
195
196    /// Protect nodes from being gc collected, for each node do following steps:
197    ///
198    /// 1. if node is local or root, do nothing
199    /// 1. if has current scope, add node to current scope
200    /// 1. can't protect, returns false
201    pub fn protect_nodes(&mut self, nodes: &[NonNull<GcHead>]) {
202        self.protect_nodes_iter(nodes.iter().copied());
203    }
204
205    /// Update memory usage with rollup to parent partitions
206    pub(crate) fn update_mem_use(&mut self, id: GcPartitionId, delta: i32) -> usize {
207        if id.is_null() {
208            return 0;
209        }
210
211        if let Some(par) = self.partitions.get_mut(&id) {
212            if delta >= 0 {
213                let d = delta as usize;
214                par.memory_used += d;
215                self.total_memory_used += d;
216            } else {
217                let d = (-delta) as usize;
218                debug_assert!(par.memory_used >= d);
219                debug_assert!(self.total_memory_used >= d);
220                par.memory_used -= d;
221                self.total_memory_used -= d;
222            }
223            par.memory_used
224        } else {
225            0
226        }
227    }
228
229    #[inline(always)]
230    pub const fn memory_used(&self) -> usize {
231        self.total_memory_used
232    }
233}
234
235#[cfg(test)]
236mod heap_tests {
237    use crate::{GcRef, GcTraceCtx, trace::GcTrace};
238
239    use super::*;
240
241    #[derive(Debug)]
242    struct Node {
243        next: Option<GcRef<Node>>,
244        value: i32,
245    }
246
247    impl GcTrace for Node {
248        fn trace(&self, tr: &mut GcTraceCtx) {
249            if let Some(next) = self.next {
250                tr.add(next);
251            }
252        }
253    }
254
255    crate::gc_type_register! {
256        Node, drop_pass = 0;
257    }
258
259    #[test]
260    fn test_heap_with_context_alloc_and_cleanup() {
261        let mut heap = GcHeap::new(&GC_TYPE_REGISTRY);
262        let partition_id = heap.create_partition();
263
264        let head = heap.with_new_scope(partition_id, |ctx| {
265            let node: GcRef<Node> = ctx
266                .alloc_local(Node {
267                    next: None,
268                    value: 1,
269                })
270                .unwrap();
271            ctx.flush();
272            node.head_ptr
273        });
274
275        while !heap.mark(partition_id, 64) {}
276        let removed_after = heap.sweep(partition_id, GcHeap::DUMMY_DISPOSE_CALLBACK);
277        assert!(removed_after > 0);
278    }
279}
gc_lite/heap.rs

gc_lite/
heap.rs
