Fix memory bugs.

src/lib.rs

@@ -1,9 +1,12 @@
 #![feature(new_range_api)]
+#![feature(coerce_unsized)]
+#![feature(unsize)]
 #![allow(warnings)]
 extern crate core;
 
 pub mod ast;
 pub mod diagnostic;
+pub mod ir;
 pub mod module;
 pub mod name_analysis;
 pub mod object_file;
@@ -11,4 +14,3 @@ pub mod parser;
 pub mod std_core;
 pub mod util;
 pub mod vm;
-pub mod ir;

src/vm/mem/mod.rs

@@ -1,90 +1,302 @@
+use std::any::TypeId;
 use std::cell::RefCell;
-use std::marker::PhantomData;
-use std::ops::Mul;
+use std::fmt::{Debug, Formatter};
+use std::marker::Unsize;
+use std::ops::{CoerceUnsized, Deref, DerefMut, Mul};
+use std::sync::OnceLock;
 
-#[derive(Debug)]
-pub struct Gc<T> {
-    inner: *mut GcInner,
-    data: *const T,
-    _phantom: PhantomData<T>,
+pub trait Trace {
+    fn add_grays(&self, grays: &mut Vec<GcAny>);
+
+    fn dynamic_size(&self) -> usize {
+        size_of_val(&self)
+    }
 }
 
-#[derive(Debug)]
-pub struct GcInner {
-    color: GcColor,
-    size: usize,
-}
-
-impl GcInner {
-    pub fn new(size: usize) -> Self {
-        Self {
-            color: GcColor::White,
-            size,
+impl Trace for RefCell<DvmObject> {
+    fn add_grays(&self, grays: &mut Vec<GcAny>) {
+        for field in self.borrow().fields() {
+            match field {
+                DvmValue::Object(dvm_object) => grays.push(dvm_object.as_any()),
+                DvmValue::NativeObject(native_object) => grays.push(native_object.clone()),
+                _ => {}
+            }
         }
     }
 }
 
-#[derive(PartialEq, Eq, Copy, Clone, Debug)]
+impl Trace for String {
+    fn add_grays(&self, grays: &mut Vec<GcAny>) {
+        // no-op
+    }
+
+    fn dynamic_size(&self) -> usize {
+        self.len()
+    }
+}
+
+#[derive(Clone)]
+pub struct Gc<T: Trace> {
+    inner: *mut GcInner<T>,
+}
+
+impl<T: Trace> Gc<T> {
+    pub fn new(data: T) -> Self {
+        let inner = Box::into_raw(Box::new(GcInner::new(data)));
+        Self { inner }
+    }
+
+    fn from_inner(inner: *mut GcInner<T>) -> Self {
+        Self { inner }
+    }
+
+    pub fn borrow(&self) -> GcRef<T> {
+        unsafe {
+            (*self.inner).borrow_count += 1;
+            GcRef::new(self.inner)
+        }
+    }
+}
+
+impl<T: Trace + 'static> Gc<T> {
+    pub fn as_any(&self) -> GcAny {
+        GcAny::new::<T>(self.inner)
+    }
+}
+
+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>,
+}
+
+impl<T: Trace> GcRef<T> {
+    fn new(inner: *mut GcInner<T>) -> Self {
+        Self { inner }
+    }
+}
+
+impl<T: Trace> Deref for GcRef<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { &(*self.inner).data }
+    }
+}
+
+impl<T: Trace> DerefMut for GcRef<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { &mut (*self.inner).data }
+    }
+}
+
+impl<T: Trace> Drop for GcRef<T> {
+    fn drop(&mut self) {
+        unsafe {
+            (*self.inner).borrow_count -= 1;
+        }
+    }
+}
+
+struct GcInner<T: Trace> {
+    allocated_size: usize,
+    color: GcColor,
+    borrow_count: usize,
+    next: Option<GcAny>,
+    data: T,
+}
+
+impl<T: Trace> GcInner<T> {
+    fn new(data: T) -> Self {
+        Self {
+            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)]
 pub enum GcColor {
     White,
     Black,
 }
 
-impl<T> Gc<T> {
-    pub fn new(data: T) -> Self {
-        let inner = Box::into_raw(Box::new(GcInner::new(size_of::<T>())));
-        let boxed_data = Box::into_raw(Box::new(data)) as *const T;
+struct GcAnyMeta {
+    trace: unsafe fn(*mut (), &mut Vec<GcAny>),
+    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,
+}
+
+#[derive(Clone)]
+pub struct GcAny {
+    inner: *mut (), // *mut GcInner<T>
+    meta: &'static GcAnyMeta,
+}
+
+impl GcAny {
+    pub fn new<T: Trace + 'static>(gc_inner: *mut GcInner<T>) -> Self {
+        unsafe fn trace_impl<T: Trace>(gc_inner: *mut (), grays: &mut Vec<GcAny>) {
+            (*(gc_inner as *mut GcInner<T>)).data.add_grays(grays);
+        }
+
+        unsafe fn allocated_size_impl<T: Trace>(gc_inner: *mut ()) -> usize {
+            (*(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 {
+            (*(gc_inner as *mut GcInner<T>)).borrow_count
+        }
+
+        unsafe fn color_impl<T: Trace>(gc_inner: *mut ()) -> GcColor {
+            (*(gc_inner as *mut GcInner<T>)).color
+        }
+
+        unsafe fn set_color_impl<T: Trace>(gc_inner: *mut (), color: GcColor) {
+            (*(gc_inner as *mut GcInner<T>)).color = color;
+        }
+
+        unsafe fn next_impl<T: Trace>(gc_inner: *mut ()) -> Option<GcAny> {
+            (*(gc_inner as *mut GcInner<T>)).next.clone()
+        }
+
+        unsafe fn set_next_impl<T: Trace>(gc_inner: *mut (), next: 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>));
+        }
+
+        static VTABLE: OnceLock<GcAnyMeta> = OnceLock::new();
+        let v_table = VTABLE.get_or_init(|| GcAnyMeta {
+            trace: trace_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>(),
+        });
         Self {
-            inner,
-            data: boxed_data,
-            _phantom: PhantomData,
+            meta: v_table,
+            inner: gc_inner as *mut (),
         }
     }
 
-    pub fn set_color(&mut self, color: GcColor) {
+    fn trace(&self, grays: &mut Vec<GcAny>) {
         unsafe {
-            (*self.inner).color = color;
+            (self.meta.trace)(self.inner, grays);
         }
     }
 
-    pub fn color(&self) -> GcColor {
-        unsafe { (*self.inner).color }
+    fn allocated_size(&self) -> usize {
+        unsafe { (self.meta.allocated_size)(self.inner) }
     }
 
-    pub fn size(&self) -> usize {
-        unsafe { (*self.inner).size }
+    fn dynamic_size(&self) -> usize {
+        unsafe { (self.meta.dynamic_size)(self.inner) }
     }
 
-    pub fn data(&self) -> &T {
-        unsafe { &*self.data }
+    fn borrow_count(&self) -> usize {
+        unsafe { (self.meta.borrow_count)(self.inner) }
     }
 
-    pub fn dealloc(&mut self) {
+    fn color(&self) -> GcColor {
+        unsafe { (self.meta.color)(self.inner) }
+    }
+
+    fn set_color(&self, color: GcColor) {
         unsafe {
-            drop(Box::from_raw(self.inner));
-            drop(Box::from_raw(self.data as *mut T));
+            (self.meta.set_color)(self.inner, color);
+        }
+    }
+
+    fn next(&self) -> Option<GcAny> {
+        unsafe { (self.meta.next)(self.inner) }
+    }
+
+    fn set_next(&self, next: GcAny) {
+        unsafe { (self.meta.set_next)(self.inner, next) }
+    }
+
+    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");
+            }
+            (self.meta.dealloc)(self.inner)
+        }
+    }
+
+    pub fn is<T: 'static>(&self) -> bool {
+        self.meta.type_id == TypeId::of::<T>()
+    }
+
+    pub fn downcast<T: Trace + 'static>(&self) -> Option<Gc<T>> {
+        if self.is::<T>() {
+            unsafe { Some(Gc::from_inner(self.inner as *mut GcInner<T>)) }
+        } else {
+            None
         }
     }
 }
 
-impl<T> Clone for Gc<T> {
-    fn clone(&self) -> Self {
-        Self {
-            inner: self.inner,
-            data: self.data,
-            _phantom: PhantomData,
-        }
+impl Debug for GcAny {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "GcAny@{:p}", self.inner)
     }
 }
 
-#[derive(Debug)]
+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>>),
+    NativeObject(GcAny),
     Nil,
 }
 
-#[derive(Debug)]
 struct DvmObject {
     fields: Vec<DvmValue>,
 }
@@ -103,103 +315,207 @@ impl DvmObject {
     }
 }
 
-fn collect_garbage(
-    next_gc: &mut usize,
-    current_size: &mut usize,
-    stack: &Vec<DvmValue>,
-    gcs: &mut Vec<Gc<RefCell<DvmObject>>>,
-) {
-    let mut gray_stack: Vec<Gc<RefCell<DvmObject>>> = vec![];
+pub struct GcHeap {
+    head: Option<GcAny>,
+    tail: Option<GcAny>,
+    next_gc: usize,
+    current_size: usize,
+}
+
+impl GcHeap {
+    pub fn new() -> Self {
+        Self {
+            head: None,
+            tail: None,
+            next_gc: 1024 * 1024,
+            current_size: 0,
+        }
+    }
+
+    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.replace(as_any.clone()).unwrap();
+            tail.set_next(as_any);
+        } else {
+            self.head = Some(as_any.clone());
+            self.tail = Some(as_any);
+        }
+        gc
+    }
+
+    pub fn head(&self) -> Option<GcAny> {
+        self.head.clone()
+    }
+
+    pub fn set_head(&mut self, head: GcAny) {
+        self.head = Some(head.clone());
+        if self.tail.is_none() {
+            self.tail = Some(head);
+        }
+    }
+
+    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 {
+        self.current_size
+    }
+
+    pub fn set_current_size(&mut self, size: usize) {
+        self.current_size = size;
+    }
+
+    pub fn subtract_current_size(&mut self, size: usize) {
+        self.current_size -= size;
+    }
+
+    pub fn next_gc(&self) -> usize {
+        self.next_gc
+    }
+
+    pub fn set_next_gc(&mut self, next_gc: usize) {
+        self.next_gc = next_gc;
+    }
+}
+
+impl Drop for GcHeap {
+    fn drop(&mut self) {
+        let mut current = self.head().take();
+        while let Some(gc_any) = &mut current {
+            let next = gc_any.next();
+            unsafe { gc_any.dealloc(); }
+            current = next;
+        }
+    }
+}
+
+fn collect_garbage(stack: &Vec<DvmValue>, heap: &mut GcHeap) {
+    if heap.current_size() < heap.next_gc() {
+        return;
+    }
+
+    let mut gray_stack: Vec<GcAny> = vec![];
 
     for stack_value in stack {
         match stack_value {
             DvmValue::Object(gc) => {
-                gray_stack.push(gc.clone());
+                gray_stack.push(gc.as_any());
+            }
+            DvmValue::NativeObject(gc_any) => {
+                gray_stack.push(gc_any.clone());
             }
             _ => {}
         }
     }
 
     while let Some(mut current) = gray_stack.pop() {
-        for field in current.data().borrow().fields() {
-            match field {
-                DvmValue::Object(object) => {
-                    if object.color() == GcColor::White {
-                        gray_stack.push(object.clone());
-                    }
-                }
-                _ => {}
-            }
-        }
+        current.trace(&mut gray_stack);
         current.set_color(GcColor::Black);
     }
 
     let mut collected_size: usize = 0;
 
-    for i in (0..gcs.len()).rev() {
-        if gcs[i].color() == GcColor::White {
-            let mut gc = gcs.remove(i);
-            let size = gc.size();
-            *current_size -= size;
-            collected_size += size;
-            gc.dealloc();
+    let mut previous: Option<GcAny> = None;
+    let mut current = heap.head();
+    while let Some(mut current_gc) = current.take() {
+        if let GcColor::White = current_gc.color() {
+            // Must delete from linked-list and dealloc
+            // Case 0: Head, next None
+            // Case 1a: Head, next Some
+            // Case 1b: Non-head, next Some
+            // 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
+            }
+            heap.subtract_current_size(current_gc.allocated_size());
+            collected_size += current_gc.dynamic_size();
+            unsafe { current_gc.dealloc(); }
         } else {
-            gcs[i].set_color(GcColor::White);
+            current_gc.set_color(GcColor::White);
+            previous = Some(current_gc.clone());
+            current = current_gc.next();
         }
     }
 
-    *next_gc = std::cmp::max(1024 * 1024, current_size.mul(2));
-    println!(
-        "collected_size: {}, current_size: {}, next_gc: {}",
-        collected_size, current_size, next_gc
-    );
-}
-
-fn maybe_collect_garbage(
-    next_gc: &mut usize,
-    current_size: &mut usize,
-    stack: &Vec<DvmValue>,
-    gcs: &mut Vec<Gc<RefCell<DvmObject>>>,
-) {
-    if current_size > next_gc {
-        collect_garbage(next_gc, current_size, stack, gcs);
+    if let Ok(val) = std::env::var("DVM_MEM_DEBUG") {
+        if val == "1" {
+            println!("GC: collected_size: {}", collected_size);
+        }
     }
-}
 
-fn alloc_dvm_object(
-    next_gc: &mut usize,
-    current_size: &mut usize,
-    stack: &Vec<DvmValue>,
-    gcs: &mut Vec<Gc<RefCell<DvmObject>>>,
-) -> Gc<RefCell<DvmObject>> {
-    maybe_collect_garbage(next_gc, current_size, stack, gcs);
-    let mut object = DvmObject::new();
-    object.fields_mut().push(DvmValue::Nil);
-    let gc = Gc::new(RefCell::new(object));
-    *current_size += gc.size();
-    gcs.push(gc.clone());
-    gc
-}
-
-fn set_field(obj: &mut Gc<RefCell<DvmObject>>, field_index: usize, value: DvmValue) {
-    let mut binding = obj.data().borrow_mut();
-    binding.fields_mut()[field_index] = value;
+    heap.set_next_gc(std::cmp::max(1024 * 1024, heap.current_size()));
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::vm::mem::{alloc_dvm_object, collect_garbage, DvmObject, DvmValue, Gc};
-    use std::cell::RefCell;
+    use crate::vm::mem::{DvmValue, GcHeap};
 
     #[test]
-    fn test() {
+    fn simple_allocate() {
         let mut stack: Vec<DvmValue> = vec![];
-        let mut heap: Vec<Gc<RefCell<DvmObject>>> = vec![];
-        let mut next_gc: usize = 1024 * 1024;
-        let mut current_size: usize = 0;
-        for i in 0..100_000_000 {
-            let mut o = alloc_dvm_object(&mut next_gc, &mut current_size, &stack, &mut heap);
+        let mut heap = GcHeap::new();
+        for i in 0..100_000 {
+            let o = heap.alloc(&stack, String::new());
+            o.borrow()
+                .push_str("hello world! long string! woot woot woot owo");
         }
-        collect_garbage(&mut next_gc, &mut current_size, &stack, &mut heap);
+    }
+
+    #[test]
+    fn downcast_as_string() {
+        let stack = vec![];
+        let mut heap = GcHeap::new();
+        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"));
     }
 }