From a9707efa5fddd12e90b3b8b9f8ab8ab6b3ee5a88 Mon Sep 17 00:00:00 2001
From: Yash Karandikar <nerdstep710@gmail.com>
Date: Sun, 19 Sep 2021 14:12:41 -0500
Subject: [PATCH] Add a heap

---
 src/allocator.rs | 24 ++++++++++++++++++++++++
 src/main.rs      | 21 +++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/src/allocator.rs b/src/allocator.rs
index c2146b9..16c4e0c 100644
--- a/src/allocator.rs
+++ b/src/allocator.rs
@@ -1,4 +1,5 @@
 use alloc::alloc::{GlobalAlloc, Layout};
+use x86_64::{VirtAddr, structures::paging::{FrameAllocator, Mapper, Page, PageTableFlags, Size4KiB, mapper::MapToError}};
 use core::ptr::null_mut;
 
 pub struct Dummy;
@@ -15,3 +16,26 @@ unsafe impl GlobalAlloc for Dummy {
 
 #[global_allocator]
 static ALLOCATOR: Dummy = Dummy;
+
+pub const HEAP_START: usize = 0x_4444_4444_0000;
+pub const HEAP_SIZE: usize = 100 * 1024; // 100 KiB
+
+pub fn init_heap(mapper: &mut impl Mapper<Size4KiB>, frame_allocator: &mut impl FrameAllocator<Size4KiB>) -> Result<(), MapToError<Size4KiB>> {
+    let page_range = {
+        let heap_start = VirtAddr::new(HEAP_START as u64);
+        let heap_end = heap_start + HEAP_SIZE - 1u64;
+        let heap_start_page = Page::containing_address(heap_start);
+        let heap_end_page = Page::containing_address(heap_end);
+        Page::range_inclusive(heap_start_page, heap_end_page)
+    };
+
+    for page in page_range {
+        let frame = frame_allocator.allocate_frame().ok_or(MapToError::FrameAllocationFailed)?;
+        let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;
+        unsafe {
+            mapper.map_to(page, frame, flags, frame_allocator)?.flush()
+        }
+    }
+
+    Ok(())
+}
diff --git a/src/main.rs b/src/main.rs
index b35cfe5..147c2a8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -38,6 +38,8 @@ fn init() {
 #[no_mangle]
 pub extern "C" fn _start(boot_info: &'static BootInfo) {
     use crate::vga_buffer::{change_color, Color};
+    use memory::BootInfoFrameAllocator;
+    use x86_64::VirtAddr;
 
     init();
     print!("[ ");
@@ -46,6 +48,25 @@ pub extern "C" fn _start(boot_info: &'static BootInfo) {
     change_color(Color::White, Color::Black);
     println!(" ] Initialized GDT and interrupts");
 
+    let phys_mem_offset = VirtAddr::new(boot_info.physical_memory_offset);
+    let mut mapper = unsafe { memory::init(phys_mem_offset) };
+    let mut frame_allocator = unsafe {
+        BootInfoFrameAllocator::init(&boot_info.memory_map)
+    };
+    print!("[ ");
+    change_color(Color::Green, Color::Black);
+    print!("OK");
+    change_color(Color::White, Color::Black);
+    println!(" ] Initialized Mapper and Frame allocator");
+
+    allocator::init_heap(&mut mapper, &mut frame_allocator).expect("Heap initialization failed");
+    print!("[ ");
+    change_color(Color::Green, Color::Black);
+    print!("OK");
+    change_color(Color::White, Color::Black);
+    println!(" ] Initialized heap");
+
+    println!();
     print!("Welcome to ");
     change_color(Color::Blue, Color::Black);
     println!("KarxOS!");