move more stuff and fix heap

kernel/src/arch/riscv64/init.rs

@@ -2,7 +2,7 @@ use core::{arch::asm, ops::Range};
 
 use crate::{
     arch::{asm, csr, interrupts, paging, wait},
-    fdt::FDT,
+    dev::fdt::FDT,
     start,
 };
 
@@ -70,5 +70,8 @@ pub unsafe fn init() -> ! {
         end: raw_mem_range.end(),
     };
     to_supervisor();
-    start(heap_mem, fdt)
+    start(crate::StartInfo {
+        mem_range: heap_mem,
+        dt: fdt,
+    })
 }

kernel/src/dev/fdt.rs

@@ -1,15 +1,19 @@
 // garbage .1% finished FDT implementation
 
+use alloc::format;
+
 use crate::{
     println,
-    util::{bits::{u32_from_bytes, Be}, lazy::LazyConst},
+    util::bits::{u32_from_bytes, Be},
 };
 use core::{
-    fmt::Debug, mem::{size_of, transmute}, ops::Range, slice
+    fmt::Debug,
+    mem::{size_of, transmute},
+    ops::Range,
+    slice,
 };
 
 const MAGIC: u32 = 0xd00dfeed;
-pub static DT: LazyConst<FDT> = LazyConst::new();
 
 #[repr(u32)]
 #[derive(Clone, Copy, Debug)]
@@ -20,7 +24,9 @@ enum Token {
     Nop,
     End,
 }
+
 const TOKEN_SIZE: usize = size_of::<Token>();
+
 impl Token {
     pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
         let val = u32_from_bytes(bytes)?.to_be();
@@ -63,6 +69,15 @@ pub struct Prop {
     pub data: &'static [u8],
 }
 
+impl Debug for Prop {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("Prop")
+            .field("name", &self.name)
+            .field("data_len", &self.data)
+            .finish()
+    }
+}
+
 impl Prop {
     pub fn full_len(&self) -> usize {
         return PROP_SIZE + self.data.len();
@@ -213,6 +228,16 @@ pub struct Node {
     pub props: &'static [u8],
 }
 
+impl Debug for Node {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        let props: alloc::vec::Vec<_> = self.into_iter().map(|p| format!("{:?}", p)).collect();
+        f.debug_struct("Node")
+            .field("name", &self.name)
+            .field("props", &props)
+            .finish()
+    }
+}
+
 impl IntoIterator for &Node {
     type Item = Prop;
     type IntoIter = PropIter;

kernel/src/dev/mod.rs

@@ -1 +1,2 @@
 pub mod uart;
+pub mod fdt;

kernel/src/main.rs

@@ -9,13 +9,12 @@
 
 use crate::mem::alloc::ALLOCATOR;
 use core::ops::Range;
-use fdt::{DT, FDT};
+use dev::fdt::FDT;
 
 extern crate alloc;
 
 pub mod arch;
 pub mod dev;
-pub mod fdt;
 pub mod log;
 pub mod mem;
 pub mod qemu;
@@ -23,20 +22,31 @@ pub mod qemu;
 mod test;
 pub mod util;
 
-pub fn start(heap_mem: Range<*mut u8>, fdt: FDT) -> ! {
+pub struct StartInfo {
-    DT.init(fdt);
+    mem_range: Range<*mut u8>,
-    unsafe {
+    dt: FDT
-        ALLOCATOR.init(heap_mem);
 }
+
+pub fn start(info: StartInfo) -> ! {
     #[cfg(test)]
+    {
+        // un... un bro momento..
+        test::init(info);
         test_main();
+    }
     #[cfg(not(test))]
-    main();
+    main(info);
     qemu::exit(0)
 }
 
-pub fn main() {
+pub fn main(info: StartInfo) {
     println!("we out here vibin");
+    unsafe {
+        ALLOCATOR.init(&info.mem_range);
+    }
+    for dev in info.dt.into_iter() {
+        println!("{:?}", dev);
+    }
 }
 
 #[panic_handler]

kernel/src/mem/alloc.rs

@@ -14,9 +14,12 @@ impl Allocator {
     pub const fn empty() -> Self {
         Self(Mutex::new(Heap::empty()))
     }
-    pub unsafe fn init(&self, range: Range<*mut u8>) {
+    pub unsafe fn init(&self, range: &Range<*mut u8>) {
         self.0.lock().init(range);
     }
+    pub unsafe fn reset(&self, range: &Range<*mut u8>) {
+        self.0.lock().reset(range);
+    }
     pub fn print(&self) {
         self.0.lock().print();
     }

kernel/src/mem/heap.rs

@@ -15,11 +15,12 @@ use crate::println;
 const ALIGN: usize = 0b1000;
 const ALIGN_MASK: usize = !(ALIGN - 1);
 
+#[repr(C)]
 pub struct BlockInfo(usize);
 
 impl BlockInfo {
-    pub const fn new(prev_used: bool, size: usize) -> Self {
+    pub const fn new(size: usize) -> Self {
-        Self(prev_used as usize | size)
+        Self(true as usize | size)
     }
     pub fn prev_used(&self) -> bool {
         self.0 & 1 == 1
@@ -37,6 +38,7 @@ impl BlockInfo {
 
 pub type BlockPointer = *mut BlockInfo;
 
+#[repr(C)]
 pub struct FreeBlockInfo {
     info: BlockInfo,
     prev: FreePointer,
@@ -62,6 +64,8 @@ pub struct Heap {
     head: FreeBlockInfo,
     start: *mut u8,
     end: *mut u8,
+    #[cfg(debug_assertions)]
+    pub debug: bool,
 }
 
 impl Heap {
@@ -74,17 +78,24 @@ impl Heap {
             },
             start: null_mut(),
             end: null_mut(),
+            #[cfg(debug_assertions)]
+            debug: false,
         }
     }
 
-    pub unsafe fn init(&mut self, range: Range<*mut u8>) {
+    pub unsafe fn reset(&mut self, range: &Range<*mut u8>) {
+        *self = Self::empty();
+        self.init(range);
+    }
+
+    pub unsafe fn init(&mut self, range: &Range<*mut u8>) {
         let head = self.head();
         let first = range.start as FreePointer;
         let size = range.end as usize - range.start as usize;
         create_free(
             first,
             FreeBlockInfo {
-                info: BlockInfo::new(true, size),
+                info: BlockInfo::new(size),
                 next: head,
                 prev: head,
             },
@@ -108,38 +119,41 @@ impl Heap {
             let free_size = (*free).info.size();
             // free block found
             if free_size >= size {
+                #[cfg(debug_assertions)]
+                {
+                    if self.debug {
+                        println!(
+                            "-------- \x1b[92malloc\x1b[0m: {:?}; size 0x{:x}",
+                            free, size
+                        );
+                    }
+                }
                 // deal with leftover space
                 let leftover = free_size - size;
                 if leftover < FREE_SIZE {
+                    // consume if not enough space for another free block
                     size = free_size;
                     let mut next_used = free.byte_add(size) as BlockPointer;
                     if next_used as *mut u8 == self.end {
                         next_used = &mut self.head.info;
                     }
                     (*next_used).set_prev_used();
-                    let prev = (*free).prev;
+                    remove_free(free);
-                    let next = (*free).next;
-                    (*prev).next = next;
-                    (*next).prev = prev;
                 } else {
+                    // otherwise create another free
                     let new_free = free.byte_add(size);
-                    let prev = (*free).prev;
+                    self.insert_free((*free).prev, (*free).next, new_free, leftover);
-                    let next = (*free).next;
-                    create_free(
-                        new_free,
-                        FreeBlockInfo {
-                            info: BlockInfo::new(true, leftover),
-                            prev,
-                            next,
-                        },
-                    );
-                    (*prev).next = new_free;
-                    (*next).prev = new_free;
                 }
                 // create block
                 let used = free as BlockPointer;
-                (*used) = BlockInfo::new(true, size);
+                (*used) = BlockInfo::new(size);
                 let data = used.byte_add(USED_SIZE) as *mut u8;
+                #[cfg(debug_assertions)]
+                {
+                    if self.debug {
+                        self.print();
+                    }
+                }
                 return data;
             }
         }
@@ -148,43 +162,79 @@ impl Heap {
 
     pub unsafe fn dealloc(&mut self, ptr: *mut u8, _: core::alloc::Layout) {
         let used = ptr.byte_sub(USED_SIZE) as BlockPointer;
+        #[cfg(debug_assertions)]
+        {
+            if self.debug {
+                println!("-------- \x1b[91mdealloc\x1b[0m: {:?}", used);
+            }
+        }
         let mut size = (*used).size();
-        let old_size = size;
         let mut addr = used as FreePointer;
-        let mut prev = self.head();
+        let next = self.next(used);
-        let mut next = self.head.next;
+        // merge with behind if free
         if !(*used).prev_used() {
-            let prev_free = *(used.byte_sub(PTR_SIZE) as *mut FreePointer);
+            let prev = *(used.byte_sub(PTR_SIZE) as *mut FreePointer);
-            addr = prev_free;
+            addr = prev;
-            size += (*prev_free).info.size();
+            size += (*prev).info.size();
-            prev = (*prev_free).prev;
+            remove_free(prev);
-            next = (*prev_free).next;
         }
-        let mut n_block = used.byte_add(old_size);
+        // merge with after if free
-        if n_block as *mut u8 != self.end {
+        if !self.is_end(next) && self.is_free(next) {
-            let mut nn_block = n_block.byte_add((*n_block).size());
+            size += (*next).size();
-            if nn_block as *mut u8 == self.end {
+            let ahead = next as FreePointer;
-                nn_block = &mut self.head.info;
+            remove_free(ahead);
         }
-            if !(*nn_block).prev_used() {
+        // create & insert
-                size += (*n_block).size();
+        let head = self.head();
-                next = (*next).next;
+        self.insert_free(head, self.head.next, addr, size);
+        (*next).unset_prev_used();
+
+        #[cfg(debug_assertions)]
+        {
+            if self.debug {
+                self.print();
             }
         }
+    }
+
+    unsafe fn insert_free(
+        &mut self,
+        prev: FreePointer,
+        next: FreePointer,
+        addr: FreePointer,
+        size: usize,
+    ) {
         create_free(
             addr,
             FreeBlockInfo {
-                info: BlockInfo::new(true, size),
+                info: BlockInfo::new(size),
                 prev,
                 next,
             },
         );
         (*prev).next = addr;
         (*next).prev = addr;
-        if n_block as *mut u8 == self.end {
-            n_block = &mut self.head.info;
     }
-        (*n_block).unset_prev_used();
+
+    unsafe fn next(&mut self, block: BlockPointer) -> BlockPointer {
+        let mut next = block.byte_add((*block).size());
+        // head is "end" of list
+        if next as *mut u8 == self.end {
+            next = &mut self.head.info;
+        }
+        next
+    }
+
+    unsafe fn is_end(&self, block: BlockPointer) -> bool {
+        block as *mut u8 == self.end
+    }
+
+    unsafe fn is_used(&mut self, block: BlockPointer) -> bool {
+        (*self.next(block)).prev_used()
+    }
+
+    unsafe fn is_free(&mut self, block: BlockPointer) -> bool {
+        !self.is_used(block)
     }
 
     pub fn iter_free(&mut self) -> FreeBlockIter {
@@ -206,21 +256,35 @@ impl Heap {
 
     pub fn print(&mut self) {
         unsafe {
-            println!("heap: {:?} -> {:?}", self.start, self.end);
+            println!("heap: {:?} -> {:?}", self.start, self.end,);
+            let ptyp = if self.head.prev_used() {
+                "used"
+            } else {
+                "free"
+            };
+            println!(" - {:?}: head, prev is {}", self.head(), ptyp);
+            println!(
+                "   L prev: {:?}, next: {:?}",
+                self.head.prev, self.head.next
+            );
             for block in self.iter_block() {
                 let size = (*block).size();
                 let mut n_block = block.byte_add(size);
                 if n_block as *mut u8 == self.end {
                     n_block = &mut self.head.info;
                 }
-                let used = if (*n_block).prev_used() {
+                let used = (*n_block).prev_used();
-                    "used"
+                let typ = if used { "used" } else { "free" };
-                } else {
+                let ptyp = if (*block).prev_used() { "used" } else { "free" };
-                    "free"
+                println!(
-                };
+                    " - {:?}: {}, prev is {}, size 0x{:x}",
-                println!(" - {:?}: {}, size 0x{:x}", block, used, size);
+                    block, typ, ptyp, size
+                );
+                if !used {
+                    let block = block as FreePointer;
+                    println!("   L prev: {:?}, next: {:?}", (*block).prev, (*block).next);
+                }
             }
-            println!();
         }
     }
 }
@@ -232,6 +296,13 @@ unsafe fn create_free(addr: FreePointer, info: FreeBlockInfo) {
     *end = addr;
 }
 
+unsafe fn remove_free(block: FreePointer) {
+    let next = (*block).next;
+    let prev = (*block).prev;
+    (*next).prev = prev;
+    (*prev).next = next;
+}
+
 pub struct FreeBlockIter {
     end: FreePointer,
     prev: FreePointer,

kernel/src/test/mem.rs

@@ -1,7 +1,12 @@
-use crate::ALLOCATOR;
+use crate::{StartInfo, ALLOCATOR};
 
 #[test_case]
-fn memory1() {
+fn init(_: &StartInfo) {
+    check_empty();
+}
+
+#[test_case]
+fn vec(_: &StartInfo) {
     let mut test = alloc::vec![1, 2, 3];
     assert_eq!(test[0], 1);
     assert_eq!(test[1], 2);
@@ -11,7 +16,7 @@ fn memory1() {
     assert_eq!(test[1], 2);
     assert_eq!(test[2], 3);
     assert_eq!(test[3], 3);
-    // check allocator
+    check_used(1);
 
     let test2 = alloc::vec![-1, -2, -3, -4];
     assert_eq!(test2[0], -1);
@@ -23,15 +28,16 @@ fn memory1() {
     assert_eq!(test[1], 2);
     assert_eq!(test[2], 3);
     assert_eq!(test[3], 3);
-    // check allocator
+    check_used(2);
 
     drop(test2);
+    check_used(1);
     drop(test);
-    // check allocator
+    check_empty();
 }
 
 #[test_case]
-fn memory2() {
+fn vec_vec(_: &StartInfo) {
     let mut test = alloc::vec::Vec::new();
     for i in 0..4 {
         let n = i * 4;
@@ -44,20 +50,43 @@ fn memory2() {
         assert_eq!(test[i][2], n + 2);
         assert_eq!(test[i][3], n + 3);
     }
-    // check allocator
+    check_used(5);
 
     drop(test);
-    // check allocator
+    check_empty();
 }
 
 #[test_case]
-fn memory_reuse() {
+fn reuse(info: &StartInfo) {
-    for _ in 0..1000 {
+    let size = info.mem_range.end as usize - info.mem_range.start as usize;
-        let _: alloc::vec::Vec<i32> = alloc::vec::Vec::with_capacity(10_000_0000);
+    for _ in 0..10 {
+        let _: alloc::vec::Vec<u8> = alloc::vec::Vec::with_capacity(size / 2);
     }
-    // check allocator
+    check_empty();
 }
 
-pub fn check_alloc_empty() {
+pub fn check_empty() {
-    ALLOCATOR.heap();
+    let mut free = 0;
+    for _ in ALLOCATOR.heap().iter_free() {
+        free += 1;
+    }
+    let mut total = 0;
+    for _ in ALLOCATOR.heap().iter_block() {
+        total += 1;
+    }
+    let used = total - free;
+    assert_eq!((free, used), (1, 0), "(free, used)");
+}
+
+pub fn check_used(num: usize) {
+    let mut free = 0;
+    for _ in ALLOCATOR.heap().iter_free() {
+        free += 1;
+    }
+    let mut total = 0;
+    for _ in ALLOCATOR.heap().iter_block() {
+        total += 1;
+    }
+    let used = total - free;
+    assert_eq!(used, num);
 }

kernel/src/test/mod.rs

@@ -1,4 +1,6 @@
-use crate::{print, println, qemu};
+use core::mem::MaybeUninit;
+
+use crate::{mem::alloc::ALLOCATOR, print, println, qemu, StartInfo};
 
 mod mem;
 
@@ -8,30 +10,58 @@ static mut TESTS: &[&dyn Testable] = &[];
 static mut TEST: usize = 0;
 static mut FAILED: usize = 0;
 
+static mut START_INFO: MaybeUninit<StartInfo> = MaybeUninit::uninit();
+
 pub trait Testable {
-    fn run(&self) -> ();
+    fn run(&self, info: &StartInfo) -> ();
 }
 
-impl<T: Fn()> Testable for T {
+// bruh...
-    fn run(&self) {
+
-        print!("test {}... ", core::any::type_name::<T>());
+// impl<T: Fn()> Testable for T {
-        self();
+//     fn run(&self, _: &StartInfo) {
+//         // TODO: very temp solution that fails if names are changed lmao
+//         const START: &str = "kernel::test::";
+//         let name = &core::any::type_name::<T>()[START.len()..];
+//         print!("test {}... ", name);
+//         self();
+//         println!("\x1b[92mok\x1b[0m");
+//     }
+// }
+
+impl<T: Fn(&StartInfo)> Testable for T {
+    fn run(&self, info: &StartInfo) {
+        // TODO: very temp solution that fails if names are changed lmao
+        const START: &str = "kernel::test::";
+        let name = &core::any::type_name::<T>()[START.len()..];
+        print!("test {}... ", name);
+        self(info);
         println!("\x1b[92mok\x1b[0m");
     }
 }
 
+pub fn init(info: StartInfo) {
+    unsafe {
+        *START_INFO.as_mut_ptr() = info;
+    }
+}
+
 pub fn test_runner(tests: &[&dyn Testable]) {
-    unsafe { TESTS = core::mem::transmute(tests) };
+    unsafe {
+        TESTS = core::mem::transmute(tests);
+    }
     println!("Running {} tests", tests.len());
     run_tests();
 }
 
 pub fn run_tests() -> ! {
+    let info = unsafe {START_INFO.assume_init_ref()};
     unsafe {
         for i in TEST..TESTS.len() {
             let test = TESTS[i];
             TEST += 1;
-            test.run();
+            prepare(info);
+            test.run(info);
         }
         print!(
             "results: {}. {} passed; {} failed",
@@ -48,6 +78,12 @@ pub fn run_tests() -> ! {
     qemu::exit(0)
 }
 
+fn prepare(info: &StartInfo) {
+    unsafe {
+        ALLOCATOR.reset(&info.mem_range);
+    }
+}
+
 pub fn test_panic(info: &core::panic::PanicInfo) -> ! {
     println!("\x1b[91mFAILED\x1b[0m");
     println!("\x1b[93m{}\x1b[0m", info);

kernel/src/util/lazy.rs

@@ -1,40 +1,39 @@
-use core::{mem::zeroed, ops::Deref};
+use core::{mem::MaybeUninit, ops::Deref};
 
 pub struct LazyConst<T> {
-    #[cfg(not(debug_assertions))]
     value: T,
     #[cfg(debug_assertions)]
-    value: Option<T>,
+    state: u8,
 }
 
 impl<T> Deref for LazyConst<T> {
     type Target = T;
 
     fn deref(&self) -> &Self::Target {
-        #[cfg(not(debug_assertions))]
-        {
-            &self.value
-        }
         #[cfg(debug_assertions)]
-        {
+        if self.state == 0 {
-            self.value.as_ref().expect("Value was not assigned")
+            panic!("Lazy const for {} not assigned", core::any::type_name::<T>());
         }
+        &self.value
     }
 }
 
 impl<T> LazyConst<T> {
     pub const fn new() -> Self {
-        unsafe { Self { value: zeroed() } }
+        unsafe {
+            Self {
+                value: MaybeUninit::zeroed().assume_init(),
+                #[cfg(debug_assertions)]
+                state: 0,
+            }
+        }
     }
 
-    pub fn init(&self, value: T) {
+    pub unsafe fn init(&self, value: T) {
-        #[cfg(not(debug_assertions))]
+        as_mut(self).value = value;
-        unsafe {
-            as_mut(self).value = value
-        }
         #[cfg(debug_assertions)]
-        unsafe {
+        {
-            as_mut(self).value = Some(value)
+            as_mut(self).state = 1;
         }
     }
 }