From d38d16154c5686a39fcacc51134731c5388aff3e Mon Sep 17 00:00:00 2001
From: Haobo Gu <haobogu@outlook.com>
Date: Fri, 27 Oct 2023 23:14:32 +0800
Subject: [PATCH] feat(core): add heap, use it for eeprom write buf

Signed-off-by: Haobo Gu <haobogu@outlook.com>
---
 Cargo.lock          | 17 +++++++++++++++++
 rmk/Cargo.toml      |  1 +
 rmk/src/eeprom.rs   |  7 +++++--
 rmk/src/keyboard.rs | 15 +++++++++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 83a667ed..139cb84b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -182,6 +182,16 @@ version = "1.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07"
 
+[[package]]
+name = "embedded-alloc"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8931e47e33c5d3194fbcf9cc82df0919193bd2fa40008f388eb1d28fd9c9ea6b"
+dependencies = [
+ "critical-section",
+ "linked_list_allocator",
+]
+
 [[package]]
 name = "embedded-dma"
 version = "0.2.0"
@@ -328,6 +338,12 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "linked_list_allocator"
+version = "0.10.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9afa463f5405ee81cdb9cc2baf37e08ec7e4c8209442b5d72c04cfb2cd6e6286"
+
 [[package]]
 name = "log"
 version = "0.4.20"
@@ -520,6 +536,7 @@ dependencies = [
  "byteorder",
  "cortex-m",
  "cortex-m-rt",
+ "embedded-alloc",
  "embedded-hal",
  "embedded-storage",
  "log",
diff --git a/rmk/Cargo.toml b/rmk/Cargo.toml
index 6da6b897..4d912777 100644
--- a/rmk/Cargo.toml
+++ b/rmk/Cargo.toml
@@ -21,6 +21,7 @@ usb-device = "0.2.9"
 usbd-hid = "0.6.1"
 packed_struct = { version = "0.10.1", default-features = false }
 byteorder = { version = "1.4", default-features = false }
+embedded-alloc = "0.5.0"
 
 [features]
 default = ["col2row"]
diff --git a/rmk/src/eeprom.rs b/rmk/src/eeprom.rs
index d83069d1..431b1073 100644
--- a/rmk/src/eeprom.rs
+++ b/rmk/src/eeprom.rs
@@ -1,8 +1,11 @@
 pub mod eeconfig;
 pub mod eekeymap;
 
+extern crate alloc;
+
 use self::eeconfig::EEPROM_MAGIC;
 use crate::{action::KeyAction, keymap::KeyMapConfig};
+use alloc::vec;
 use core::sync::atomic::{AtomicBool, Ordering::SeqCst};
 use embedded_storage::nor_flash::NorFlash;
 use log::{debug, error, info, warn};
@@ -65,6 +68,7 @@ pub struct Eeprom<F: NorFlash, const EEPROM_SIZE: usize> {
     ///  EEPROM_SIZE should be at least 1008(keymap) + 15(eeconfig) + 100(macro)
     keymap_config: KeyMapConfig,
 }
+
 impl<F: NorFlash, const EEPROM_SIZE: usize> Eeprom<F, EEPROM_SIZE> {
     pub fn new<const ROW: usize, const COL: usize, const NUM_LAYER: usize>(
         storage: F,
@@ -223,8 +227,7 @@ impl<F: NorFlash, const EEPROM_SIZE: usize> Eeprom<F, EEPROM_SIZE> {
         }
 
         let bytes = record.to_bytes();
-        // TODO: Use a buf whose length equals storage_write_size
-        let mut buf = [0xFF_u8; 16];
+        let mut buf = vec![0xFF_u8; self.storage_config.page_size as usize];
         buf[..bytes.len()].copy_from_slice(&bytes);
         debug!(
             "EEPROM write storage at 0x{:X}: {:02X?} ",
diff --git a/rmk/src/keyboard.rs b/rmk/src/keyboard.rs
index 25c62837..2f65776c 100644
--- a/rmk/src/keyboard.rs
+++ b/rmk/src/keyboard.rs
@@ -8,6 +8,7 @@ use crate::{
     via::{descriptor::ViaReport, process::process_via_packet},
 };
 use core::convert::Infallible;
+use embedded_alloc::Heap;
 use embedded_hal::digital::v2::{InputPin, OutputPin};
 use embedded_storage::nor_flash::NorFlash;
 use log::debug;
@@ -15,6 +16,9 @@ use rtic_monotonics::systick::*;
 use usb_device::class_prelude::UsbBus;
 use usbd_hid::descriptor::{KeyboardReport, MediaKeyboardReport, SystemControlReport};
 
+#[global_allocator]
+static HEAP: Heap = Heap::empty();
+
 pub struct Keyboard<
     In: InputPin,
     Out: OutputPin,
@@ -75,6 +79,17 @@ impl<
         eeprom_storage_config: EepromStorageConfig,
         mut keymap: [[[KeyAction; COL]; ROW]; NUM_LAYER],
     ) -> Self {
+        // Initialize the allocator at the very beginning of the initialization of the keyboard
+        {
+            use core::mem::MaybeUninit;
+            // 1KB heap size
+            const HEAP_SIZE: usize = 1024;
+            // Check page_size and heap size
+            assert!((eeprom_storage_config.page_size as usize) < HEAP_SIZE);
+            static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
+            unsafe { HEAP.init(HEAP_MEM.as_ptr() as usize, HEAP_SIZE) }
+        }
+
         let eeprom = match storage {
             Some(s) => {
                 let e = Eeprom::new(s, eeprom_storage_config, &keymap);