sodium_rust/impl_/
gc_node.rs

1use std::cell::Cell;
2use std::collections::HashSet;
3use std::fmt::Write as _;
4use std::sync::Arc;
5use std::sync::Mutex;
6use std::sync::RwLock;
7
8use log::trace;
9
10pub type Tracer<'a> = dyn FnMut(&GcNode) + 'a;
11
12pub type Trace = dyn Fn(&mut Tracer) + Send + Sync;
13
14#[derive(PartialEq, Eq, Clone, Copy)]
15enum Color {
16    Black,
17    Gray,
18    Purple,
19    White,
20}
21
22#[derive(Clone)]
23pub struct GcNode {
24    id: u32,
25    name: String,
26    gc_ctx: GcCtx,
27    data: Arc<GcNodeData>,
28}
29
30struct GcNodeData {
31    freed: Cell<bool>,
32    ref_count: Cell<u32>,
33    ref_count_adj: Cell<u32>,
34    visited: Cell<bool>,
35    color: Cell<Color>,
36    buffered: Cell<bool>,
37    deconstructor: RwLock<Box<dyn Fn() + Send + Sync>>,
38    trace: RwLock<Box<Trace>>,
39}
40
41unsafe impl Send for GcNodeData {}
42unsafe impl Sync for GcNodeData {}
43
44#[derive(Clone)]
45pub struct GcCtx {
46    data: Arc<Mutex<GcCtxData>>,
47}
48
49struct GcCtxData {
50    next_id: u32,
51    roots: Vec<GcNode>,
52    to_be_freed: Vec<GcNode>,
53}
54
55impl Default for GcCtx {
56    fn default() -> GcCtx {
57        GcCtx::new()
58    }
59}
60
61impl GcCtx {
62    pub fn new() -> GcCtx {
63        GcCtx {
64            data: Arc::new(Mutex::new(GcCtxData {
65                next_id: 0,
66                roots: Vec::new(),
67                to_be_freed: Vec::new(),
68            })),
69        }
70    }
71
72    fn with_data<R, K: FnOnce(&mut GcCtxData) -> R>(&self, k: K) -> R {
73        let mut l = self.data.lock();
74        let data = l.as_mut().unwrap();
75        k(data)
76    }
77
78    pub fn make_id(&self) -> u32 {
79        self.with_data(|data: &mut GcCtxData| {
80            let id = data.next_id;
81            data.next_id += 1;
82            id
83        })
84    }
85
86    pub fn add_possible_root(&self, node: GcNode) {
87        self.with_data(|data: &mut GcCtxData| data.roots.push(node));
88    }
89
90    pub fn collect_cycles(&self) {
91        loop {
92            trace!("start: collect_cycles");
93            self.mark_roots();
94            self.scan_roots();
95            self.collect_roots();
96            trace!("end: collect_cycles");
97            let bail = self.with_data(|data: &mut GcCtxData| {
98                data.roots.is_empty() && data.to_be_freed.is_empty()
99            });
100            if bail {
101                break;
102            }
103        }
104    }
105
106    fn mark_roots(&self) {
107        trace!("start: mark_roots");
108        let mut old_roots: Vec<GcNode> = Vec::new();
109        self.with_data(|data: &mut GcCtxData| std::mem::swap(&mut old_roots, &mut data.roots));
110        self.display_graph(&old_roots);
111        let mut new_roots: Vec<GcNode> = Vec::new();
112        for root in &old_roots {
113            self.reset_ref_count_adj_step_1_of_2(root);
114        }
115        for root in &old_roots {
116            self.reset_ref_count_adj_step_2_of_2(root);
117        }
118        for root in old_roots {
119            let color = root.data.color.get();
120            if color == Color::Purple {
121                self.mark_gray(&root);
122                new_roots.push(root);
123            } else {
124                root.data.buffered.set(false);
125                if root.data.color.get() == Color::Black
126                    && root.data.ref_count.get() == 0
127                    && !root.data.freed.get()
128                {
129                    self.with_data(|data: &mut GcCtxData| data.to_be_freed.push(root));
130                }
131            }
132        }
133        self.with_data(|data: &mut GcCtxData| std::mem::swap(&mut new_roots, &mut data.roots));
134        trace!("end: mark_roots");
135    }
136
137    fn display_graph(&self, roots: &[GcNode]) {
138        let mut stack = Vec::new();
139        let mut visited: HashSet<*const GcNodeData> = HashSet::new();
140        let mut show_names_for = Vec::new();
141        for root in roots {
142            stack.push(root.clone());
143        }
144        trace!("-- start of graph drawing --");
145        loop {
146            let next_op = stack.pop();
147            if next_op.is_none() {
148                break;
149            }
150            let next = next_op.unwrap();
151            {
152                let next_ptr: &GcNodeData = &next.data;
153                let next_ptr: *const GcNodeData = next_ptr;
154                if visited.contains(&next_ptr) {
155                    continue;
156                }
157                visited.insert(next_ptr);
158            }
159            show_names_for.push(next.clone());
160            let mut line: String =
161                format!("id {}, ref_count {}: ", next.id, next.data.ref_count.get());
162            let mut first: bool = true;
163            next.trace(|t| {
164                if first {
165                    first = false;
166                } else {
167                    line.push(',');
168                }
169                write!(line, "{}", t.id).ok();
170                stack.push(t.clone());
171            });
172            trace!("{}", line);
173        }
174        trace!("node names:");
175        for next in show_names_for {
176            trace!("{}: {}", next.id, next.name);
177        }
178        trace!("-- end of graph drawing --");
179    }
180
181    fn mark_gray(&self, s: &GcNode) {
182        if s.data.color.get() == Color::Gray {
183            return;
184        }
185        s.data.color.set(Color::Gray);
186
187        s.trace(&mut |t: &GcNode| {
188            trace!("mark_gray: gc node {} dec ref count", t.id);
189            t.data.ref_count_adj.set(t.data.ref_count_adj.get() + 1);
190            if t.data.ref_count_adj.get() > t.data.ref_count.get() {
191                panic!("ref count adj was larger than ref count for node {} ({}) (ref adj {}) (ref cnt {})", t.id, t.name, t.data.ref_count_adj.get(), t.data.ref_count.get());
192            }
193            self.mark_gray(t);
194        });
195    }
196
197    fn scan_roots(&self) {
198        trace!("start: scan_roots");
199        let mut roots = Vec::new();
200        self.with_data(|data: &mut GcCtxData| std::mem::swap(&mut roots, &mut data.roots));
201        for root in &roots {
202            self.scan(root);
203        }
204        for root in &roots {
205            self.reset_ref_count_adj_step_1_of_2(root);
206        }
207        for root in &roots {
208            self.reset_ref_count_adj_step_2_of_2(root);
209        }
210        self.with_data(|data: &mut GcCtxData| std::mem::swap(&mut roots, &mut data.roots));
211        trace!("end: scan_roots");
212    }
213
214    fn scan(&self, s: &GcNode) {
215        if s.data.color.get() != Color::Gray {
216            return;
217        }
218        if s.data.ref_count_adj.get() == s.data.ref_count.get() {
219            s.data.color.set(Color::White);
220            trace!("scan: gc node {} became white", s.id);
221            s.trace(|t| {
222                self.scan(t);
223            });
224        } else {
225            self.scan_black(s);
226        }
227    }
228
229    fn reset_ref_count_adj_step_1_of_2(&self, s: &GcNode) {
230        if s.data.visited.get() {
231            return;
232        }
233        s.data.visited.set(true);
234        s.data.ref_count_adj.set(0);
235        s.trace(|t| {
236            self.reset_ref_count_adj_step_1_of_2(t);
237        });
238    }
239
240    fn reset_ref_count_adj_step_2_of_2(&self, s: &GcNode) {
241        if !s.data.visited.get() {
242            return;
243        }
244        s.data.visited.set(false);
245        s.trace(|t| {
246            self.reset_ref_count_adj_step_2_of_2(t);
247        });
248    }
249
250    fn scan_black(&self, s: &GcNode) {
251        s.data.color.set(Color::Black);
252        trace!("scan: gc node {} became black", s.id);
253        let this = self.clone();
254        s.trace(|t| {
255            if t.data.color.get() != Color::Black {
256                this.scan_black(t);
257            }
258        });
259    }
260
261    fn collect_roots(&self) {
262        let mut white = Vec::new();
263        let mut roots = Vec::new();
264        self.with_data(|data: &mut GcCtxData| roots.append(&mut data.roots));
265        for root in &roots {
266            root.data.buffered.set(false);
267            self.collect_white(root, &mut white);
268        }
269        for i in &white {
270            if !i.data.freed.get() {
271                trace!("collect_roots: freeing white node {} ({})", i.id, i.name);
272                i.free();
273                self.with_data(|data: &mut GcCtxData| {
274                    data.roots.retain(|root: &GcNode| root.id != i.id)
275                });
276            }
277        }
278        let mut to_be_freed = Vec::new();
279        self.with_data(|data: &mut GcCtxData| to_be_freed.append(&mut data.to_be_freed));
280        for i in &to_be_freed {
281            if !i.data.freed.get() {
282                trace!(
283                    "collect_roots: freeing to_be_freed node {} ({})",
284                    i.id,
285                    i.name
286                );
287                i.free();
288                self.with_data(|data: &mut GcCtxData| {
289                    data.roots.retain(|root: &GcNode| root.id != i.id)
290                });
291            }
292        }
293        for i in white {
294            if i.ref_count() != 0 {
295                panic!(
296                    "freed node ref count did not drop to zero for node {} ({})",
297                    i.id, i.name
298                );
299            }
300        }
301        for i in to_be_freed {
302            if i.ref_count() != 0 {
303                panic!(
304                    "freed node ref count did not drop to zero for node {} ({})",
305                    i.id, i.name
306                );
307            }
308        }
309    }
310
311    fn collect_white(&self, s: &GcNode, white: &mut Vec<GcNode>) {
312        if s.data.color.get() == Color::White
313        /*&& !s.data.buffered.get()*/
314        {
315            s.data.color.set(Color::Black);
316            s.trace(|t| {
317                self.collect_white(t, white);
318            });
319            trace!("collect_white: gc node {} added to white list", s.id);
320            white.push(s.clone());
321        }
322    }
323}
324
325impl GcNode {
326    pub fn new<
327        NAME: ToString,
328        DECONSTRUCTOR: 'static + Fn() + Send + Sync,
329        TRACE: 'static + Fn(&mut Tracer) + Send + Sync,
330    >(
331        gc_ctx: &GcCtx,
332        name: NAME,
333        deconstructor: DECONSTRUCTOR,
334        trace: TRACE,
335    ) -> GcNode {
336        GcNode {
337            id: gc_ctx.make_id(),
338            name: name.to_string(),
339            gc_ctx: gc_ctx.clone(),
340            data: Arc::new(GcNodeData {
341                freed: Cell::new(false),
342                ref_count: Cell::new(1),
343                ref_count_adj: Cell::new(0),
344                visited: Cell::new(false),
345                color: Cell::new(Color::Black),
346                buffered: Cell::new(false),
347                deconstructor: RwLock::new(Box::new(deconstructor)),
348                trace: RwLock::new(Box::new(trace)),
349            }),
350        }
351    }
352
353    pub fn ref_count(&self) -> u32 {
354        self.data.ref_count.get()
355    }
356
357    pub fn inc_ref_if_alive(&self) -> bool {
358        if self.ref_count() != 0 && !self.data.freed.get() {
359            self.data.ref_count.set(self.data.ref_count.get() + 1);
360            self.data.color.set(Color::Black);
361            true
362        } else {
363            false
364        }
365    }
366
367    pub fn inc_ref(&self) {
368        if self.data.freed.get() {
369            panic!("gc_node {} inc_ref on freed node ({})", self.id, self.name);
370        }
371        self.data.ref_count.set(self.data.ref_count.get() + 1);
372        self.data.color.set(Color::Black);
373    }
374
375    pub fn dec_ref(&self) {
376        if self.data.ref_count.get() == 0 {
377            return;
378        }
379        self.data.ref_count.set(self.data.ref_count.get() - 1);
380        if self.data.ref_count.get() == 0 {
381            self.release();
382        } else {
383            self.possible_root();
384        }
385    }
386
387    pub fn release(&self) {
388        self.data.color.set(Color::Black);
389        if !self.data.buffered.get() {
390            trace!("release: freeing gc_node {} ({})", self.id, self.name);
391            self.free();
392        }
393    }
394
395    pub fn possible_root(&self) {
396        if self.data.color.get() != Color::Purple {
397            self.data.color.set(Color::Purple);
398            if !self.data.buffered.get() {
399                self.data.buffered.set(true);
400                self.gc_ctx.add_possible_root(self.clone());
401            }
402        }
403    }
404
405    pub fn free(&self) {
406        self.data.freed.set(true);
407        let mut tmp: Box<dyn Fn() + Send + Sync + 'static> = Box::new(|| {});
408        {
409            let mut deconstructor = self.data.deconstructor.write().unwrap();
410            std::mem::swap(&mut *deconstructor, &mut tmp);
411        }
412        tmp();
413        let mut trace = self.data.trace.write().unwrap();
414        *trace = Box::new(|_tracer: &mut Tracer| {});
415    }
416
417    pub fn trace<TRACER: FnMut(&GcNode)>(&self, mut tracer: TRACER) {
418        let trace = self.data.trace.read().unwrap();
419        trace(&mut tracer);
420    }
421}