1 files changed, 32 insertions, 30 deletions

diff --git a/src/devices/memory.rs b/src/devices/memory.rs
index 1d33f64..56125d2 100644
--- a/src/devices/memory.rs
+++ b/src/devices/memory.rs
@@ -5,8 +5,11 @@ fn new_blank_page() -> [u8; 65536] { [0; 65536] }
 
 
 pub struct MemoryDevice {
-    pages: Vec<Page>,
+    // Hard limit on the number of pages which can be allocated
     pub page_limit: u16,
+    // Pages which have actually been allocated
+    pub pages: Vec<Page>,
+
     pub page_1: u16,
     pub address_1: u16,
     pub page_2: u16,
@@ -17,8 +20,9 @@ pub struct MemoryDevice {
 impl MemoryDevice {
     pub fn new() -> Self {
         Self {
-            pages: Vec::new(),
             page_limit: 0,
+            pages: Vec::new(),
+
             page_1: 0,
             address_1: 0,
             page_2: 0,
@@ -27,31 +31,21 @@ impl MemoryDevice {
         }
     }
 
-    pub fn page_count(&self) -> u16 {
-        self.pages.len() as u16
-    }
-
     pub fn copy(&mut self) {
-        // Return if at least one page is out-of-bounds
-        if self.page_1 as usize >= self.pages.len()
-        || self.page_2 as usize >= self.pages.len() {
-            return
-        }
-        let p1 = self.page_1 as usize;
-        let p2 = self.page_2 as usize;
-        let a1 = self.address_1;
-        let a2 = self.address_2;
-
-        for i in 0..=self.copy_length {
-            let byte = self.pages[p2][a2.wrapping_add(i) as usize];
-            self.pages[p1][a1.wrapping_add(i) as usize] = byte;
-        }
-    }
+        let count = std::cmp::max(self.page_1, self.page_2) as usize + 1;
+        if count <= self.page_limit as usize {
+            if self.pages.len() < count {
+                self.pages.resize_with(count, new_blank_page);
+            }
+            let p1 = self.page_1 as usize;
+            let p2 = self.page_2 as usize;
+            let a1 = self.address_1;
+            let a2 = self.address_2;
 
-    pub fn expand_memory(&mut self) {
-        let size = std::cmp::max(self.page_1, self.page_2) as usize + 1;
-        if size < self.page_limit as usize {
-            self.pages.resize_with(size, new_blank_page);
+            for i in 0..=self.copy_length {
+                let byte = self.pages[p2][a2.wrapping_add(i) as usize];
+                self.pages[p1][a1.wrapping_add(i) as usize] = byte;
+            }
         }
     }
 
@@ -75,15 +69,23 @@ impl MemoryDevice {
 
     pub fn write_to_head_1(&mut self, byte: u8) {
         let address = then_inc!(self.address_1);
-        if let Some(page) = self.pages.get_mut(self.page_1 as usize) {
-            page[address] = byte;
+        let page = self.page_1 as usize;
+        if self.page_limit as usize > page {
+            if self.pages.len() <= page {
+                self.pages.resize_with(page + 1, new_blank_page);
+            }
+            self.pages[page][address] = byte;
         }
     }
 
     pub fn write_to_head_2(&mut self, byte: u8) {
-        let address = then_inc!(self.address_2);
-        if let Some(page) = self.pages.get_mut(self.page_2 as usize) {
-            page[address] = byte;
+        let address = then_inc!(self.address_1);
+        let page = self.page_2 as usize;
+        if self.page_limit as usize > page {
+            if self.pages.len() <= page {
+                self.pages.resize_with(page + 1, new_blank_page);
+            }
+            self.pages[page][address] = byte;
         }
     }
 }