Cleanup heap management

kernel/context/memory.rs

@@ -21,6 +21,14 @@ impl Memory {
         memory
     }
 
+    pub fn start_address(&self) -> VirtualAddress {
+        self.start
+    }
+
+    pub fn size(&self) -> usize {
+        self.size
+    }
+
     pub fn pages(&self) -> PageIter {
         let start_page = Page::containing_address(self.start);
         let end_page = Page::containing_address(VirtualAddress::new(self.start.get() + self.size - 1));
@@ -52,14 +60,52 @@ impl Memory {
     pub fn remap(&mut self, new_flags: EntryFlags, flush: bool) {
         let mut active_table = unsafe { ActivePageTable::new() };
 
-        self.flags = new_flags;
-
         for page in self.pages() {
-            active_table.remap(page, self.flags);
+            active_table.remap(page, new_flags);
             if flush {
                 active_table.flush(page);
             }
         }
+
+        self.flags = new_flags;
+    }
+
+    pub fn resize(&mut self, new_size: usize, flush: bool) {
+        let mut active_table = unsafe { ActivePageTable::new() };
+
+        if new_size > self.size {
+            let start_page = Page::containing_address(VirtualAddress::new(self.start.get() + self.size));
+            let end_page = Page::containing_address(VirtualAddress::new(self.start.get() + new_size - 1));
+            for page in Page::range_inclusive(start_page, end_page) {
+                //println!("Map {:X}", page.start_address().get());
+                if active_table.translate_page(page).is_none() {
+                    //println!("Not found - mapping");
+                    active_table.map(page, self.flags);
+                    if flush {
+                        active_table.flush(page);
+                    }
+                } else {
+                    //println!("Found - skipping");
+                }
+            }
+        } else if new_size < self.size {
+            let start_page = Page::containing_address(VirtualAddress::new(self.start.get() + new_size));
+            let end_page = Page::containing_address(VirtualAddress::new(self.start.get() + self.size - 1));
+            for page in Page::range_inclusive(start_page, end_page) {
+                //println!("Unmap {:X}", page.start_address().get());
+                if active_table.translate_page(page).is_some() {
+                    //println!("Found - unmapping");
+                    active_table.unmap(page);
+                    if flush {
+                        active_table.flush(page);
+                    }
+                } else {
+                    //println!("Not found - skipping");
+                }
+            }
+        }
+
+        self.size = new_size;
     }
 }
 

kernel/context/mod.rs

@@ -178,6 +178,8 @@ pub struct Context {
     pub image: Vec<memory::Memory>,
     /// User stack
     pub stack: Option<memory::Memory>,
+    /// User heap
+    pub heap: Option<memory::Memory>,
     /// The open files in the scheme
     pub files: Vec<Option<file::File>>
 }
@@ -193,6 +195,7 @@ impl Context {
             kstack: None,
             image: Vec::new(),
             stack: None,
+            heap: None,
             files: Vec::new()
         }
     }