Fix memory bugs.

src/vm/mem/mod.rs

@@ -1,11 +1,16 @@
 use std::any::TypeId;
 use std::cell::RefCell;
+use std::fmt::{Debug, Formatter};
 use std::marker::Unsize;
 use std::ops::{CoerceUnsized, Deref, DerefMut, Mul};
 use std::sync::OnceLock;
 
 pub trait Trace {
     fn add_grays(&self, grays: &mut Vec<GcAny>);
+
+    fn dynamic_size(&self) -> usize {
+        size_of_val(&self)
+    }
 }
 
 impl Trace for RefCell<DvmObject> {
@@ -24,6 +29,10 @@ impl Trace for String {
     fn add_grays(&self, grays: &mut Vec<GcAny>) {
         // no-op
     }
+
+    fn dynamic_size(&self) -> usize {
+        self.len()
+    }
 }
 
 #[derive(Clone)]
@@ -33,8 +42,7 @@ pub struct Gc<T: Trace> {
 
 impl<T: Trace> Gc<T> {
     pub fn new(data: T) -> Self {
-        let size = size_of_val(&data);
+        let inner = Box::into_raw(Box::new(GcInner::new(data)));
-        let inner = Box::into_raw(Box::new(GcInner::new(data, size)));
         Self { inner }
     }
 
@@ -56,6 +64,12 @@ impl<T: Trace + 'static> Gc<T> {
     }
 }
 
+impl<T: Trace> Debug for Gc<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Gc@{:p}", self.inner)
+    }
+}
+
 pub struct GcRef<T: Trace> {
     inner: *mut GcInner<T>,
 }
@@ -89,7 +103,7 @@ impl<T: Trace> Drop for GcRef<T> {
 }
 
 struct GcInner<T: Trace> {
-    size: usize,
+    allocated_size: usize,
     color: GcColor,
     borrow_count: usize,
     next: Option<GcAny>,
@@ -97,15 +111,19 @@ struct GcInner<T: Trace> {
 }
 
 impl<T: Trace> GcInner<T> {
-    fn new(data: T, size: usize) -> Self {
+    fn new(data: T) -> Self {
         Self {
-            size,
+            allocated_size: size_of::<T>(),
             color: GcColor::White,
             borrow_count: 0,
             next: None,
             data,
         }
     }
+
+    fn dynamic_size(&self) -> usize {
+        self.data.dynamic_size()
+    }
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -116,12 +134,14 @@ pub enum GcColor {
 
 struct GcAnyMeta {
     trace: unsafe fn(*mut (), &mut Vec<GcAny>),
-    size: unsafe fn(*mut ()) -> usize,
+    allocated_size: unsafe fn(*mut ()) -> usize,
+    dynamic_size: unsafe fn(*mut ()) -> usize,
     borrow_count: unsafe fn(*mut ()) -> usize,
     color: unsafe fn(*mut ()) -> GcColor,
     set_color: unsafe fn(*mut (), GcColor),
     next: unsafe fn(*mut ()) -> Option<GcAny>,
     set_next: unsafe fn(*mut (), GcAny),
+    clear_next: unsafe fn(*mut ()),
     dealloc: unsafe fn(*mut ()),
     type_id: TypeId,
 }
@@ -138,8 +158,12 @@ impl GcAny {
             (*(gc_inner as *mut GcInner<T>)).data.add_grays(grays);
         }
 
-        unsafe fn size_impl<T: Trace>(gc_inner: *mut ()) -> usize {
+        unsafe fn allocated_size_impl<T: Trace>(gc_inner: *mut ()) -> usize {
-            (*(gc_inner as *mut GcInner<T>)).size
+            (*(gc_inner as *mut GcInner<T>)).allocated_size
+        }
+
+        unsafe fn dynamic_size_impl<T: Trace>(gc_inner: *mut ()) -> usize {
+            (*(gc_inner as *mut GcInner<T>)).dynamic_size()
         }
 
         unsafe fn borrow_count_impl<T: Trace>(gc_inner: *mut ()) -> usize {
@@ -162,6 +186,10 @@ impl GcAny {
             (*(gc_inner as *mut GcInner<T>)).next = Some(next);
         }
 
+        unsafe fn clear_next_impl<T: Trace>(gc_inner: *mut ()) {
+            (*(gc_inner as *mut GcInner<T>)).next = None;
+        }
+
         unsafe fn dealloc_impl<T: Trace>(gc_inner: *mut ()) {
             drop(Box::from_raw(gc_inner as *mut GcInner<T>));
         }
@@ -169,12 +197,14 @@ impl GcAny {
         static VTABLE: OnceLock<GcAnyMeta> = OnceLock::new();
         let v_table = VTABLE.get_or_init(|| GcAnyMeta {
             trace: trace_impl::<T>,
-            size: size_impl::<T>,
+            allocated_size: allocated_size_impl::<T>,
+            dynamic_size: dynamic_size_impl::<T>,
             borrow_count: borrow_count_impl::<T>,
             color: color_impl::<T>,
             set_color: set_color_impl::<T>,
             next: next_impl::<T>,
             set_next: set_next_impl::<T>,
+            clear_next: clear_next_impl::<T>,
             dealloc: dealloc_impl::<T>,
             type_id: TypeId::of::<T>(),
         });
@@ -190,8 +220,12 @@ impl GcAny {
         }
     }
 
-    fn size(&self) -> usize {
+    fn allocated_size(&self) -> usize {
-        unsafe { (self.meta.size)(self.inner) }
+        unsafe { (self.meta.allocated_size)(self.inner) }
+    }
+
+    fn dynamic_size(&self) -> usize {
+        unsafe { (self.meta.dynamic_size)(self.inner) }
     }
 
     fn borrow_count(&self) -> usize {
@@ -216,7 +250,11 @@ impl GcAny {
         unsafe { (self.meta.set_next)(self.inner, next) }
     }
 
-    fn dealloc(&mut self) {
+    fn clear_next(&mut self) {
+        unsafe { (self.meta.clear_next)(self.inner) }
+    }
+
+    unsafe fn dealloc(&mut self) {
         unsafe {
             if self.borrow_count() > 0 {
                 panic!("Attempt to dealloc a GcAny whose borrow count > 0");
@@ -238,6 +276,20 @@ impl GcAny {
     }
 }
 
+impl Debug for GcAny {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "GcAny@{:p}", self.inner)
+    }
+}
+
+impl PartialEq for GcAny {
+    fn eq(&self, other: &Self) -> bool {
+        self.inner == other.inner
+    }
+}
+
+impl Eq for GcAny {}
+
 enum DvmValue {
     Primitive(usize),
     Object(Gc<RefCell<DvmObject>>),
@@ -280,15 +332,18 @@ impl GcHeap {
         }
     }
 
-    pub fn alloc<T: Trace + 'static>(&mut self, t: T) -> Gc<T> {
+    pub fn alloc<T: Trace + 'static>(&mut self, stack: &Vec<DvmValue>, t: T) -> Gc<T> {
+        collect_garbage(stack, self);
         let gc = Gc::new(t);
+        let as_any = gc.as_any();
+        let size = as_any.allocated_size();
+        self.current_size += size;
         if self.tail.is_some() {
-            let tail = self.tail.take().unwrap();
+            let tail = self.tail.replace(as_any.clone()).unwrap();
-            tail.set_next(gc.as_any());
+            tail.set_next(as_any);
-            self.tail = Some(tail);
         } else {
-            self.head = Some(gc.as_any());
+            self.head = Some(as_any.clone());
-            self.tail = Some(gc.as_any());
+            self.tail = Some(as_any);
         }
         gc
     }
@@ -304,8 +359,13 @@ impl GcHeap {
         }
     }
 
-    pub fn clear_head(&mut self) {
+    pub fn set_tail(&mut self, tail: GcAny) {
+        self.tail = Some(tail);
+    }
+
+    pub fn clear_head_and_tail(&mut self) {
         self.head = None;
+        self.tail = None;
     }
 
     pub fn current_size(&self) -> usize {
@@ -334,7 +394,7 @@ impl Drop for GcHeap {
         let mut current = self.head().take();
         while let Some(gc_any) = &mut current {
             let next = gc_any.next();
-            gc_any.dealloc();
+            unsafe { gc_any.dealloc(); }
             current = next;
         }
     }
@@ -368,22 +428,53 @@ fn collect_garbage(stack: &Vec<DvmValue>, heap: &mut GcHeap) {
 
     let mut previous: Option<GcAny> = None;
     let mut current = heap.head();
-    while let Some(ref mut current_gc_erased) = current {
+    while let Some(mut current_gc) = current.take() {
-        if let GcColor::White = current_gc_erased.color() {
+        if let GcColor::White = current_gc.color() {
-            if let Some(ref previous_gc_erased) = previous {
+            // Must delete from linked-list and dealloc
-                // attach previous to next if both are Some
+            // Case 0: Head, next None
-                let next = current_gc_erased.next();
+            // Case 1a: Head, next Some
-                if let Some(next_gc_erased) = next {
+            // Case 1b: Non-head, next Some
-                    previous_gc_erased.set_next(next_gc_erased);
+            // Case 2: Tail (next None)
+
+            if current_gc == heap.head().unwrap() && current_gc.next().is_none() {
+                // Case 0
+                heap.clear_head_and_tail();
+                // no need to set current; it was already taken
+                // no need to set previous
+            } else if let Some(next_gc) = current_gc.next() {
+                if current_gc == heap.head().unwrap() {
+                    // Case 1a
+                    heap.set_head(next_gc);
+                    current = heap.head(); // current is new head
+                                           // no previous since this is the head
+                } else {
+                    // Case 1b
+                    let previous_gc = previous.as_ref().unwrap();
+                    previous_gc.set_next(next_gc.clone());
+                    current = Some(next_gc);
+                    // previous stays the same
                 }
+            } else {
+                // Case 2
+                let mut previous_gc = previous.as_mut().unwrap();
+                previous_gc.clear_next();
+                heap.set_tail(previous_gc.clone());
+                // no need to set current since it was taken
+                // previous doesn't matter since we're done
             }
-            let size = current_gc_erased.size();
+            heap.subtract_current_size(current_gc.allocated_size());
-            heap.subtract_current_size(size);
+            collected_size += current_gc.dynamic_size();
-            collected_size = size;
+            unsafe { current_gc.dealloc(); }
-            current_gc_erased.dealloc();
         } else {
-            current_gc_erased.set_color(GcColor::White);
+            current_gc.set_color(GcColor::White);
-            previous = Some(current_gc_erased.clone());
+            previous = Some(current_gc.clone());
+            current = current_gc.next();
+        }
+    }
+
+    if let Ok(val) = std::env::var("DVM_MEM_DEBUG") {
+        if val == "1" {
+            println!("GC: collected_size: {}", collected_size);
         }
     }
 
@@ -398,19 +489,33 @@ mod tests {
     fn simple_allocate() {
         let mut stack: Vec<DvmValue> = vec![];
         let mut heap = GcHeap::new();
-        for i in 0..100 {
+        for i in 0..100_000 {
-            let o = heap.alloc(String::new());
+            let o = heap.alloc(&stack, String::new());
-            o.borrow().push_str("hello");
+            o.borrow()
+                .push_str("hello world! long string! woot woot woot owo");
         }
     }
 
     #[test]
     fn downcast_as_string() {
+        let stack = vec![];
         let mut heap = GcHeap::new();
-        let t = heap.alloc(String::from("hello world"));
+        let t = heap.alloc(&stack, String::from("hello world"));
         let a = t.as_any();
         if let Some(s) = a.downcast::<String>() {
             assert_eq!(*s.borrow(), "hello world");
         }
     }
+
+    #[test]
+    fn find_memory_leak() {
+        let stack = vec![];
+        let mut heap = GcHeap::new();
+        println!("alloc s0");
+        let s0 = heap.alloc(&stack, String::from("s0"));
+        println!("alloc s1");
+        let s1 = heap.alloc(&stack, String::from("s1"));
+        println!("alloc s2");
+        let s2 = heap.alloc(&stack, String::from("s2"));
+    }
 }