Implement allocator tracing and optimize internal RAM usage

firmware/Cargo.lock

@@ -26,17 +26,17 @@ dependencies = [
  "embedded-storage-async",
  "embuild",
  "enumset",
- "esp-alloc",
+ "esp-alloc 0.9.0 (git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1)",
  "esp-backtrace",
  "esp-bootloader-esp-idf",
  "esp-hal",
  "esp-hal-bounce-buffers",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-println",
  "esp-radio",
  "esp-rtos",
  "esp-storage",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "gix",
  "hex",
  "hmac",
@@ -2044,8 +2044,23 @@ dependencies = [
  "cfg-if",
  "document-features",
  "enumset",
- "esp-config",
+ "esp-config 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "linked_list_allocator",
+ "rlsf",
+]
+
+[[package]]
+name = "esp-alloc"
+version = "0.9.0"
+source = "git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1#ee6e26f2fefa4da2168c95839bf618e1ecc22cc1"
+dependencies = [
+ "allocator-api2 0.3.1",
+ "cfg-if",
+ "document-features",
+ "enumset",
+ "esp-config 0.6.1 (git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1)",
+ "esp-sync 0.1.1 (git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1)",
  "linked_list_allocator",
  "rlsf",
 ]
@@ -2058,13 +2073,13 @@ checksum = "3318413fb566c7227387f67736cf70cd74d80a11f2bb31c7b95a9eb48d079669"
 dependencies = [
  "cfg-if",
  "document-features",
- "esp-config",
+ "esp-config 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-println",
  "heapless 0.9.2",
  "riscv",
  "semihosting",
- "xtensa-lx",
+ "xtensa-lx 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -2076,9 +2091,9 @@ dependencies = [
  "cfg-if",
  "document-features",
  "embedded-storage",
- "esp-config",
+ "esp-config 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-hal-procmacros",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-rom-sys",
  "jiff",
  "log",
@@ -2092,7 +2107,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "102871054f8dd98202177b9890cb4b71d0c6fe1f1413b7a379a8e0841fc2473c"
 dependencies = [
  "document-features",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "serde",
+ "serde_yaml",
+ "somni-expr",
+]
+
+[[package]]
+name = "esp-config"
+version = "0.6.1"
+source = "git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1#ee6e26f2fefa4da2168c95839bf618e1ecc22cc1"
+dependencies = [
+ "document-features",
+ "esp-metadata-generated 0.3.0 (git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1)",
  "serde",
  "serde_yaml",
  "somni-expr",
@@ -2125,12 +2152,12 @@ dependencies = [
  "embedded-io-async 0.6.1",
  "embedded-io-async 0.7.0",
  "enumset",
- "esp-config",
+ "esp-config 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-hal-procmacros",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-riscv-rt",
  "esp-rom-sys",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-synopsys-usb-otg",
  "esp32",
  "esp32c2",
@@ -2150,7 +2177,7 @@ dependencies = [
  "riscv",
  "strum",
  "ufmt-write",
- "xtensa-lx",
+ "xtensa-lx 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "xtensa-lx-rt",
 ]
 
@@ -2185,6 +2212,11 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9a93e39c8ad8d390d248dc7b9f4b59a873f313bf535218b8e2351356972399e3"
 
+[[package]]
+name = "esp-metadata-generated"
+version = "0.3.0"
+source = "git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1#ee6e26f2fefa4da2168c95839bf618e1ecc22cc1"
+
 [[package]]
 name = "esp-phy"
 version = "0.1.1"
@@ -2193,10 +2225,10 @@ checksum = "6b1facf348e1e251517278fc0f5dc134e95e518251f5796cfbb532ca226a29bf"
 dependencies = [
  "cfg-if",
  "document-features",
- "esp-config",
+ "esp-config 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-hal",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-wifi-sys",
 ]
 
@@ -2207,8 +2239,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5a30e6c9fbcc01c348d46706fef8131c7775ab84c254a3cd65d0cd3f6414d592"
 dependencies = [
  "document-features",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "log",
  "portable-atomic",
 ]
@@ -2227,14 +2259,14 @@ dependencies = [
  "embedded-io 0.7.1",
  "embedded-io-async 0.6.1",
  "embedded-io-async 0.7.0",
- "esp-alloc",
+ "esp-alloc 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "esp-config",
+ "esp-config 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-hal",
  "esp-hal-procmacros",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-phy",
  "esp-radio-rtos-driver",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-wifi-sys",
  "heapless 0.9.2",
  "instability",
@@ -2270,7 +2302,7 @@ checksum = "cd66cccc6dd2d13e9f33668a57717ab14a6d217180ec112e6be533de93e7ecbf"
 dependencies = [
  "cfg-if",
  "document-features",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -2286,12 +2318,12 @@ dependencies = [
  "embassy-sync 0.7.2",
  "embassy-time-driver",
  "embassy-time-queue-utils",
- "esp-config",
+ "esp-config 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-hal",
  "esp-hal-procmacros",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-radio-rtos-driver",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "portable-atomic",
 ]
 
@@ -2305,9 +2337,9 @@ dependencies = [
  "embedded-storage",
  "esp-hal",
  "esp-hal-procmacros",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "esp-rom-sys",
- "esp-sync",
+ "esp-sync 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -2320,10 +2352,24 @@ dependencies = [
  "document-features",
  "embassy-sync 0.6.2",
  "embassy-sync 0.7.2",
- "esp-metadata-generated",
+ "esp-metadata-generated 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "log",
  "riscv",
- "xtensa-lx",
+ "xtensa-lx 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
+name = "esp-sync"
+version = "0.1.1"
+source = "git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1#ee6e26f2fefa4da2168c95839bf618e1ecc22cc1"
+dependencies = [
+ "cfg-if",
+ "document-features",
+ "embassy-sync 0.6.2",
+ "embassy-sync 0.7.2",
+ "esp-metadata-generated 0.3.0 (git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1)",
+ "riscv",
+ "xtensa-lx 0.13.0 (git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1)",
 ]
 
 [[package]]
@@ -8516,6 +8562,14 @@ dependencies = [
  "critical-section",
 ]
 
+[[package]]
+name = "xtensa-lx"
+version = "0.13.0"
+source = "git+https://github.com/esp-rs/esp-hal?rev=ee6e26f2fefa4da2168c95839bf618e1ecc22cc1#ee6e26f2fefa4da2168c95839bf618e1ecc22cc1"
+dependencies = [
+ "critical-section",
+]
+
 [[package]]
 name = "xtensa-lx-rt"
 version = "0.21.0"
@@ -8523,7 +8577,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8709f037fb123fe7ff146d2bce86f9dc0dfc53045c016bfd9d703317b6502845"
 dependencies = [
  "document-features",
- "xtensa-lx",
+ "xtensa-lx 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "xtensa-lx-rt-proc-macros",
 ]
 
@@ -8848,3 +8902,11 @@ dependencies = [
  "syn 2.0.114",
  "winnow",
 ]
+
+[[patch.unused]]
+name = "i-slint-core"
+version = "1.14.1"
+
+[[patch.unused]]
+name = "slint"
+version = "1.14.1"

firmware/acid-firmware/.cargo/config.toml

@@ -14,6 +14,9 @@ rustflags = [
 
   # Required to obtain backtraces on riscv (e.g. when using the "esp-backtrace" crate.)
   # "-C", "force-frame-pointers",
+
+  # Output linker map
+  # "-C", "link-arg=-Wl,-Map=target/linker.map"
 ]
 
 [env] # These must be kept in sync with /.zed/settings.json
@@ -40,11 +43,11 @@ ACID_COMPOSE_LOCALE = "cs_CZ.UTF-8"
 build-std = ["alloc", "core"]
 
 [patch.crates-io]
-# esp-backtrace = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }
+esp-backtrace = { git = "https://github.com/Limeth/esp-hal.git", rev = "114977583886be4ed866ad7b7c6f16865148e899" }
+esp-println = { git = "https://github.com/Limeth/esp-hal.git", rev = "114977583886be4ed866ad7b7c6f16865148e899" }
 # esp-hal = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }
 # esp-storage = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }
 # esp-alloc = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }
-# esp-println = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }
 # esp-radio = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }
 # esp-rtos = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }
 # esp-bootloader-esp-idf = { git = "https://github.com/Limeth/esp-hal.git", rev = "95d8c8b046e945e41294d5577528d0a1c4b03247" }

firmware/acid-firmware/Cargo.toml

@@ -6,7 +6,7 @@ description = "Firmware for the ACID keyboard"
 edition = "2024"
 
 [features]
-default = ["usb-log", "limit-fps"]
+default = ["usb-log", "limit-fps", "no-alloc-tracing"]
 # Make RMK not to use USB
 no-usb = ["rmk/_no_usb"]
 # Let RMK use BLE
@@ -32,6 +32,8 @@ probe = ["limit-fps", "rtt-log", "no-usb", "ble"]
 format-db = []
 # Avoid entering the critical section for the whole duration of printing a message to console.
 racy-logging = []
+# Global allocator tracing proxy
+no-alloc-tracing = ["esp-alloc/global-allocator"]
 
 [dependencies]
 rmk = { version = "0.8.2", git = "https://github.com/Limeth/rmk", rev = "1661c55f5c21e7d80ea3f93255df483302c74b84", default-features = false, features = [
@@ -51,7 +53,7 @@ esp-backtrace = { version = "0.18", default-features = false, features = [
 ] }
 esp-hal = { version = "1.0", features = ["esp32s3", "unstable", "psram", "log-04"] }
 esp-storage = { version = "0.8", features = ["esp32s3"] }
-esp-alloc = { version = "0.9", features = ["nightly"] }
+esp-alloc = { version = "0.9", git = "https://github.com/esp-rs/esp-hal", rev = "ee6e26f2fefa4da2168c95839bf618e1ecc22cc1", default-features = false, features = ["esp32s3", "nightly", "compat"] }
 esp-println = { version = "0.16", features = ["esp32s3", "log-04"] }
 esp-radio = { version = "0.17", features = ["esp32s3", "unstable", "ble"], optional = true }
 esp-rtos = { version = "0.2", features = ["esp32s3", "esp-radio", "embassy"] }

firmware/acid-firmware/src/db/mod.rs

@@ -16,7 +16,7 @@ use esp_hal::rng::Trng;
 use esp_storage::FlashStorage;
 use log::{debug, info};
 
-use crate::ram::PSRAM_ALLOCATOR;
+use crate::ram::{PSRAM_ALLOCATOR, PsramAllocator};
 
 pub type PartitionAcid =
     Partition<'static, CriticalSectionRawMutex, BlockingAsync<FlashStorage<'static>>>;
@@ -27,7 +27,7 @@ struct AlignedBuf<const N: usize>(pub [u8; N]);
 
 pub struct EkvFlash<T> {
     flash: T,
-    buffer: Box<AlignedBuf<{ ekv::config::PAGE_SIZE }>, &'static esp_alloc::EspHeap>,
+    buffer: Box<AlignedBuf<{ ekv::config::PAGE_SIZE }>, PsramAllocator>,
 }
 
 impl<T> EkvFlash<T> {

firmware/acid-firmware/src/ffi/mod.rs

@@ -149,7 +149,8 @@ pub unsafe extern "C" fn __xkbc_atoi(s: *const c_char) -> c_int {
     todo!()
 }
 
-// TODO: What is this even for?
+// A pointer to an array of character attributes.
+// This is used by `isdigit()`, `isalpha()`, `isspace()`, etc.
 #[unsafe(no_mangle)]
 pub static __spre__ctype_: [c_char; 0] = [];
 

firmware/acid-firmware/src/flash.rs

@@ -10,7 +10,7 @@ use log::info;
 use rmk::storage::async_flash_wrapper;
 use static_cell::StaticCell;
 
-use crate::PSRAM_ALLOCATOR;
+use crate::{PSRAM_ALLOCATOR, ram::PsramAllocator};
 
 pub type Partition = embassy_embedded_hal::flash::partition::Partition<
     'static,
@@ -25,8 +25,7 @@ pub struct Partitions {
 
 /// Initialize the flash
 pub fn initialize(flash_peripheral: esp_hal::peripherals::FLASH<'static>) -> Partitions {
-    static PARTITION_TABLE_BUFFER: StaticCell<Vec<u8, &'static esp_alloc::EspHeap>> =
+    static PARTITION_TABLE_BUFFER: StaticCell<Vec<u8, PsramAllocator>> = StaticCell::new();
-        StaticCell::new();
     let partition_table_buffer = PARTITION_TABLE_BUFFER.init_with(|| {
         let mut buffer = Vec::<u8, _>::new_in(&PSRAM_ALLOCATOR);
         buffer.resize(

firmware/acid-firmware/src/main.rs

@@ -9,6 +9,7 @@
 #![no_std]
 #![no_main]
 #![feature(allocator_api)]
+#![feature(btreemap_alloc)]
 #![feature(macro_metavar_expr)]
 #![feature(c_variadic)]
 #![feature(c_size_t)]
@@ -391,7 +392,7 @@ async fn main_task(peripherals: MainPeripherals) {
 
     // TODO: Probably want to select! instead and re-try.
     join_all![
-        run_alloc_stats_reporter(),
+        ram::run_alloc_stats_reporter(),
         initialize_and_run_rmk_devices(peripherals.matrix),
         keyboard.run(), // Keyboard is special
         run_rmk(
@@ -418,34 +419,6 @@ async fn main_task(peripherals: MainPeripherals) {
     .await;
 }
 
-async fn run_alloc_stats_reporter() {
-    let mut psram_used_prev = 0;
-    let mut heap_used_prev = 0;
-    loop {
-        let psram_stats = PSRAM_ALLOCATOR.stats();
-        let heap_stats = esp_alloc::HEAP.stats();
-        if psram_stats.current_usage != psram_used_prev {
-            let difference = psram_stats.current_usage as isize - psram_used_prev as isize;
-            psram_used_prev = psram_stats.current_usage;
-            warn!(
-                "PSRAM usage changed: {}{}\n{psram_stats}",
-                if difference < 0 { '-' } else { '+' },
-                difference.abs()
-            );
-        }
-        if heap_stats.current_usage != heap_used_prev {
-            let difference = heap_stats.current_usage as isize - heap_used_prev as isize;
-            heap_used_prev = heap_stats.current_usage;
-            warn!(
-                "HEAP usage changed: {}{}\n{heap_stats}",
-                if difference < 0 { '-' } else { '+' },
-                difference.abs()
-            );
-        }
-        Timer::after_secs(1).await;
-    }
-}
-
 async fn initialize_and_run_rmk_devices(matrix_peripherals: MatrixPeripherals) {
     // Initialize the matrix and keyboard
     const I2C_ADDR_MATRIX_LEFT: I2cAddress = I2cAddress::SevenBit(0b0100000);

firmware/acid-firmware/src/proxy.rs

@@ -11,7 +11,7 @@ use alloc::vec::Vec;
 use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
 use embassy_sync::channel::Channel;
 use embassy_time::Instant;
-use esp_alloc::{EspHeap, MemoryCapability};
+use esp_alloc::MemoryCapability;
 use log::{debug, error, info, warn};
 use rmk::descriptor::KeyboardReport;
 use rmk::hid::Report;
@@ -19,6 +19,7 @@ use rmk::{heapless, join_all};
 use slint::platform::Key;
 use xkbcommon::xkb::{self, FeedResult, KeyDirection, Keysym, ModMask, Status};
 
+use crate::ram::PsramAllocator;
 use crate::util::{DurationExt, get_file_name};
 use crate::{KEYBOARD_REPORT_PROXY, PSRAM_ALLOCATOR};
 
@@ -401,8 +402,8 @@ struct KeysymEntry {
     mask: xkb::ModMask,
 }
 
-type KeysymEntries = Vec<KeysymEntry, &'static EspHeap>;
+type KeysymEntries = Vec<KeysymEntry, PsramAllocator>;
-type KeysymMap = BTreeMap<Keysym, KeysymEntries, &'static EspHeap>;
+type KeysymMap = BTreeMap<Keysym, KeysymEntries, PsramAllocator>;
 
 /// Based on https://github.com/xkbcommon/libxkbcommon/blob/6c67e3d41d3215ab1edd4406de215c7bf1f20c74/tools/how-to-type.c#L434
 fn lookup_keysym_entries(
@@ -427,7 +428,7 @@ pub fn string_to_hid_keycodes(
     keymap: &xkb::Keymap,
     _compose_table: &xkb::compose::Table,
     string: &str,
-) -> Result<Vec<HidKeycodeWithMods, &'static EspHeap>, char> {
+) -> Result<Vec<HidKeycodeWithMods, PsramAllocator>, char> {
     #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
     struct KeycodeChoice {
         mod_mask: ModMask,

firmware/acid-firmware/src/ram.rs

@@ -1,10 +1,16 @@
+use embassy_time::Timer;
 use esp_alloc::{HeapRegion, MemoryCapability};
 use esp_hal::ram;
-use log::info;
+use log::{info, warn};
 
 // Memory allocation regions.
 // These can be debugged using `xtensa-esp32s3-elf-size -A <path-to-binary>`.
 // A panic such as `memory allocation of 3740121773 bytes failed` is caused by a heap overflow. The size is `DEEDBAAD` in hex.
+//
+// RAM usage of static variables can be diagnosed with:
+// ```sh
+// xtensa-esp32s3-elf-nm.exe -S --size-sort -t d ../target/xtensa-esp32s3-none-elf/release/acid-firmware | grep -iE ' [dbv] ' | tail -n 10 | tac
+// ```
 
 /// Total heap size
 pub const HEAP_SIZE: usize = 112 * 1024;
@@ -12,6 +18,7 @@ pub const HEAP_SIZE: usize = 112 * 1024;
 pub const STACK_SIZE_CORE_APP: usize = 80 * 1024;
 
 pub static PSRAM_ALLOCATOR: esp_alloc::EspHeap = esp_alloc::EspHeap::empty();
+pub type PsramAllocator = &'static esp_alloc::EspHeap;
 
 pub fn initialize(psram_peripheral: esp_hal::peripherals::PSRAM) {
     // Use the internal DRAM as the heap.
@@ -23,6 +30,10 @@ pub fn initialize(psram_peripheral: esp_hal::peripherals::PSRAM) {
 
     esp_alloc::heap_allocator!(#[ram(reclaimed)] size: HEAP_SIZE_RECLAIMED);
     esp_alloc::heap_allocator!(size: HEAP_SIZE - HEAP_SIZE_RECLAIMED);
+
+    #[cfg(not(feature = "no-alloc-tracing"))]
+    alloc_tracing::install_iram_allocator_proxy();
+
     info!("IRAM heap initialized!\n{}", esp_alloc::HEAP.stats());
 
     // Initialize the PSRAM allocator.
@@ -42,3 +53,318 @@ pub fn initialize(psram_peripheral: esp_hal::peripherals::PSRAM) {
         );
     }
 }
+
+pub async fn run_alloc_stats_reporter() {
+    let mut psram_used_prev = 0;
+    let mut heap_used_prev = 0;
+    loop {
+        let psram_stats = PSRAM_ALLOCATOR.stats();
+        let heap_stats = esp_alloc::HEAP.stats();
+        if psram_stats.current_usage != psram_used_prev {
+            let difference = psram_stats.current_usage as isize - psram_used_prev as isize;
+            psram_used_prev = psram_stats.current_usage;
+            warn!(
+                "PSRAM heap usage changed: {}{}\n{psram_stats}",
+                if difference < 0 { '-' } else { '+' },
+                difference.abs()
+            );
+        }
+        if heap_stats.current_usage != heap_used_prev {
+            let difference = heap_stats.current_usage as isize - heap_used_prev as isize;
+            heap_used_prev = heap_stats.current_usage;
+            warn!(
+                "IRAM heap usage changed: {}{}\n{heap_stats}",
+                if difference < 0 { '-' } else { '+' },
+                difference.abs()
+            );
+            #[cfg(not(feature = "no-alloc-tracing"))]
+            alloc_tracing::report_stats();
+        }
+        Timer::after_secs(1).await;
+    }
+}
+
+#[cfg(not(feature = "no-alloc-tracing"))]
+mod alloc_tracing {
+    use core::{
+        alloc::GlobalAlloc,
+        cell::RefCell,
+        cmp::{Ordering, Reverse},
+        fmt::Display,
+    };
+
+    use alloc::collections::btree_map::BTreeMap;
+    use embassy_sync::blocking_mutex::{Mutex, raw::CriticalSectionRawMutex};
+    use esp_alloc::EspHeap;
+    use esp_backtrace::Backtrace;
+    use esp_sync::NonReentrantMutex;
+    use log::warn;
+    use tinyvec::ArrayVec;
+
+    use crate::ram::{PSRAM_ALLOCATOR, PsramAllocator};
+
+    #[derive(Default, Clone, PartialEq, Eq, Debug)]
+    struct Stats {
+        /// The total number of allocations.
+        allocations: usize,
+        /// The number of bytes allocated in total.
+        allocated_total: usize,
+        /// The number of bytes allocated minus the number of bytes deallocated.
+        allocated_current: usize,
+    }
+
+    impl Stats {
+        fn update(&mut self, allocations_delta: isize, bytes_delta: isize, bytes_new: usize) {
+            self.allocations = (self.allocations as isize + allocations_delta) as usize;
+            self.allocated_total += bytes_new;
+            self.allocated_current = (self.allocated_current as isize + bytes_delta) as usize;
+        }
+    }
+
+    impl PartialOrd for Stats {
+        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+            Some(self.cmp(other))
+        }
+    }
+
+    impl Ord for Stats {
+        fn cmp(&self, other: &Self) -> Ordering {
+            self.allocated_current
+                .cmp(&other.allocated_current)
+                .then_with(|| self.allocated_total.cmp(&other.allocated_total).reverse())
+                .then_with(|| self.allocations.cmp(&other.allocations).reverse())
+        }
+    }
+
+    #[derive(Clone)]
+    struct BacktraceWrapper(Backtrace);
+
+    impl Display for BacktraceWrapper {
+        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+            let mut it = self.0.frames().iter();
+            if let Some(frame) = it.next() {
+                write!(f, "0x{:08x}", frame.program_counter())?;
+            }
+            while let Some(frame) = it.next() {
+                write!(f, " 0x{:08x}", frame.program_counter())?;
+            }
+            Ok(())
+        }
+    }
+
+    impl PartialEq for BacktraceWrapper {
+        fn eq(&self, other: &Self) -> bool {
+            if self.0.frames().len() != other.0.frames().len() {
+                return false;
+            }
+
+            for (lhs_frame, rhs_frame) in self.0.frames().iter().zip(other.0.frames()) {
+                if lhs_frame.program_counter() != rhs_frame.program_counter() {
+                    return false;
+                }
+            }
+
+            true
+        }
+    }
+
+    impl Eq for BacktraceWrapper {}
+
+    impl PartialOrd for BacktraceWrapper {
+        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+            Some(self.cmp(other))
+        }
+    }
+
+    impl Ord for BacktraceWrapper {
+        fn cmp(&self, other: &Self) -> Ordering {
+            let mut lhs_it = self.0.frames().iter().rev();
+            let mut rhs_it = other.0.frames().iter().rev();
+
+            loop {
+                match (lhs_it.next(), rhs_it.next()) {
+                    (Some(lhs_frame), Some(rhs_frame)) => {
+                        let ordering = lhs_frame
+                            .program_counter()
+                            .cmp(&rhs_frame.program_counter());
+
+                        if ordering != Ordering::Equal {
+                            return ordering;
+                        }
+                    }
+                    (Some(_), None) => return Ordering::Greater,
+                    (None, Some(_)) => return Ordering::Less,
+                    (None, None) => return Ordering::Equal,
+                }
+            }
+        }
+    }
+
+    #[derive(Clone)]
+    struct AllocationTracer {
+        bt_to_stats: BTreeMap<BacktraceWrapper, Stats, PsramAllocator>,
+        ptr_to_bt: BTreeMap<*mut u8, BacktraceWrapper, PsramAllocator>,
+    }
+
+    unsafe impl Send for AllocationTracer {}
+
+    impl AllocationTracer {
+        const fn new() -> Self {
+            Self {
+                bt_to_stats: BTreeMap::new_in(&PSRAM_ALLOCATOR),
+                ptr_to_bt: BTreeMap::new_in(&PSRAM_ALLOCATOR),
+            }
+        }
+
+        fn alloc(&mut self, bt: Backtrace, ptr: *mut u8, bytes: usize) {
+            let bt = BacktraceWrapper(bt.clone());
+            self.ptr_to_bt.insert(ptr, bt.clone());
+            let stats = self
+                .bt_to_stats
+                .entry(bt)
+                .or_insert_with(|| Default::default());
+            stats.update(1, bytes as isize, bytes);
+        }
+
+        fn realloc(
+            &mut self,
+            ptr_old: *mut u8,
+            ptr_new: *mut u8,
+            bytes_old: usize,
+            bytes_new: usize,
+        ) {
+            let bt = self.ptr_to_bt.remove(&ptr_old).unwrap();
+            self.ptr_to_bt.insert(ptr_new, bt.clone());
+            let stats = self.bt_to_stats.get_mut(&bt).unwrap();
+            stats.update(0, bytes_new as isize - bytes_old as isize, bytes_new);
+        }
+
+        fn dealloc(&mut self, ptr: *mut u8, bytes: usize) {
+            let bt = self.ptr_to_bt.remove(&ptr).unwrap();
+            let stats = self.bt_to_stats.get_mut(&bt).unwrap();
+            stats.update(-1, -(bytes as isize), 0);
+        }
+    }
+
+    struct TracingAllocator<T: GlobalAlloc + 'static> {
+        inner: NonReentrantMutex<Option<&'static T>>,
+        tracer: Mutex<CriticalSectionRawMutex, RefCell<AllocationTracer>>,
+    }
+
+    impl<T> TracingAllocator<T>
+    where
+        T: GlobalAlloc,
+    {
+        const fn new() -> Self {
+            Self {
+                inner: NonReentrantMutex::new(None),
+                tracer: Mutex::new(RefCell::new(AllocationTracer::new())),
+            }
+        }
+
+        fn with_inner<R>(&self, callback: impl FnOnce(&T) -> R) -> R {
+            self.inner.with(|inner| {
+                (callback)(
+                    inner
+                        .as_ref()
+                        .expect("an allocator must be installed in the global allocator proxy"),
+                )
+            })
+        }
+
+        fn update_with<R>(&self, callback: impl FnOnce(&mut AllocationTracer) -> R) -> R {
+            self.tracer
+                .lock(|tracer| (callback)(&mut tracer.borrow_mut()))
+        }
+    }
+
+    unsafe impl<T> GlobalAlloc for TracingAllocator<T>
+    where
+        T: GlobalAlloc,
+    {
+        unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
+            let bt = Backtrace::capture();
+            let ptr = self.with_inner(|inner| unsafe { inner.alloc(layout) });
+            self.update_with(|tracer| tracer.alloc(bt, ptr, layout.size()));
+            ptr
+        }
+
+        unsafe fn alloc_zeroed(&self, layout: core::alloc::Layout) -> *mut u8 {
+            let bt = Backtrace::capture();
+            let ptr = self.with_inner(|inner| unsafe { inner.alloc_zeroed(layout) });
+            self.update_with(|tracer| tracer.alloc(bt, ptr, layout.size()));
+            ptr
+        }
+
+        unsafe fn realloc(
+            &self,
+            ptr_old: *mut u8,
+            layout: core::alloc::Layout,
+            new_size: usize,
+        ) -> *mut u8 {
+            let ptr_new =
+                self.with_inner(|inner| unsafe { inner.realloc(ptr_old, layout, new_size) });
+            self.update_with(|tracer| tracer.realloc(ptr_old, ptr_new, layout.size(), new_size));
+            ptr_new
+        }
+
+        unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) {
+            self.update_with(|tracer| tracer.dealloc(ptr, layout.size()));
+            self.with_inner(|inner| unsafe { inner.dealloc(ptr, layout) });
+        }
+    }
+
+    #[global_allocator]
+    static PROXY: TracingAllocator<EspHeap> = TracingAllocator::new();
+
+    pub fn install_iram_allocator_proxy() {
+        PROXY.inner.with(|inner| {
+            *inner = Some(&esp_alloc::HEAP);
+        });
+    }
+
+    pub fn report_stats() {
+        let AllocationTracer {
+            bt_to_stats,
+            ptr_to_bt,
+        } = PROXY.tracer.lock(|tracer| tracer.borrow().clone());
+        // Wrapped in `Option` because of the `Default` requirement.
+        let mut sorted = ArrayVec::<[Option<(BacktraceWrapper, Reverse<Stats>)>; 5]>::new();
+        let bt_to_stats_len = bt_to_stats.len();
+
+        for (key, value) in bt_to_stats {
+            // Reverse ordering of stats because we are interested in the largest.
+            let value_rev = Reverse(value);
+            if let Some((_, value_last)) = sorted.last().and_then(|last| last.as_ref())
+                && &value_rev >= value_last
+                && sorted.len() >= sorted.capacity()
+            {
+                // This stat is not large enough to be inserted.
+                continue;
+            }
+            let index = match sorted.binary_search_by_key(&Some(&value_rev), |item| {
+                item.as_ref().map(|(_, current_value)| current_value)
+            }) {
+                Ok(index) => index,
+                Err(index) => index,
+            };
+            if sorted.len() >= sorted.capacity() {
+                assert!(
+                    index < sorted.len(),
+                    "the stat should be large enough to be inserted not at the end of the list"
+                );
+                let _ = sorted.pop().unwrap();
+            }
+            sorted.insert(index, Some((key, value_rev)));
+        }
+
+        warn!("Largest allocations in global allocator:");
+
+        for (index, (key, value)) in sorted.into_iter().map(Option::unwrap).enumerate() {
+            warn!("{}. {}\n{:#?}", index + 1, key, value.0);
+        }
+
+        warn!("bt_to_stats.len() = {}", bt_to_stats_len);
+        warn!("ptr_to_bt.len() = {}", ptr_to_bt.len());
+    }
+}

firmware/acid-firmware/src/ui/mod.rs

@@ -36,6 +36,7 @@ use crate::{
     },
     ffi::{alloc::__spre_free, crypto::ACTIVE_ENCRYPTED_USER_KEY},
     proxy::OUTPUT_STRING_CHANNEL,
+    ram::PsramAllocator,
     ui::{
         backend::SlintBackend,
         messages::{
@@ -130,7 +131,7 @@ pub async fn run_renderer_task(backend: SlintBackend, flash_part_acid: Partition
 
 struct State {
     window: AppWindow,
-    db: Rc<AcidDatabase>,
+    db: Rc<AcidDatabase, PsramAllocator>,
     users: SpectreUsersConfig,
     /// Currently active view.
     view: AppState,
@@ -142,20 +143,23 @@ struct State {
 }
 
 impl State {
-    async fn new(db: AcidDatabase, main: AppWindow) -> Rc<RefCell<Self>> {
+    async fn new(db: AcidDatabase, main: AppWindow) -> Rc<RefCell<Self>, PsramAllocator> {
         let users = Self::load_users(&db).await;
         let usernames = users.users.clone().map(|user| user.username);
 
-        let state = Rc::new(RefCell::new(State {
+        let state = Rc::new_in(
-            window: main.clone_strong(),
+            RefCell::new(State {
-            users,
+                window: main.clone_strong(),
-            db: Rc::new(db),
+                users,
-            view: AppState::Login,
+                db: Rc::new_in(db, &PSRAM_ALLOCATOR),
-            state_login: Default::default(),
+                view: AppState::Login,
-            state_users: Default::default(),
+                state_login: Default::default(),
-            state_user_edit: Default::default(),
+                state_users: Default::default(),
-            state_user_sites: Default::default(),
+                state_user_edit: Default::default(),
-        }));
+                state_user_sites: Default::default(),
+            }),
+            &PSRAM_ALLOCATOR,
+        );
 
         main.on_enter_view({
             let state = state.clone();
@@ -332,7 +336,10 @@ impl State {
         }
     }
 
-    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
+    fn process_callback_message(
+        state_rc: &Rc<RefCell<State>, PsramAllocator>,
+        message: CallbackMessage,
+    ) {
         let view = state_rc.borrow().view;
         match view {
             AppState::Login => StateLogin::process_callback_message(state_rc, message),
@@ -382,14 +389,21 @@ impl State {
 }
 
 trait AppViewTrait {
-    fn process_callback_message(_state_rc: &Rc<RefCell<State>>, _message: CallbackMessage) {}
+    fn process_callback_message(
+        _state_rc: &Rc<RefCell<State>, PsramAllocator>,
+        _message: CallbackMessage,
+    ) {
+    }
 }
 
 #[derive(Default)]
 struct StateLogin {}
 
 impl AppViewTrait for StateLogin {
-    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
+    fn process_callback_message(
+        state_rc: &Rc<RefCell<State>, PsramAllocator>,
+        message: CallbackMessage,
+    ) {
         let mut state = state_rc.borrow_mut();
         match message {
             CallbackMessage::Login(CallbackMessageLogin::PwAccepted {
@@ -542,7 +556,10 @@ impl AppViewTrait for StateLogin {
 struct StateUsers {}
 
 impl AppViewTrait for StateUsers {
-    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
+    fn process_callback_message(
+        state_rc: &Rc<RefCell<State>, PsramAllocator>,
+        message: CallbackMessage,
+    ) {
         let mut state = state_rc.borrow_mut();
         match message {
             CallbackMessage::Escape => {
@@ -572,7 +589,10 @@ struct StateUserEdit {
 }
 
 impl AppViewTrait for StateUserEdit {
-    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
+    fn process_callback_message(
+        state_rc: &Rc<RefCell<State>, PsramAllocator>,
+        message: CallbackMessage,
+    ) {
         let state = state_rc.clone();
         let mut state = state.borrow_mut();
         match message {
@@ -722,7 +742,10 @@ struct StateUserSites {
 }
 
 impl AppViewTrait for StateUserSites {
-    fn process_callback_message(state_rc: &Rc<RefCell<State>>, message: CallbackMessage) {
+    fn process_callback_message(
+        state_rc: &Rc<RefCell<State>, PsramAllocator>,
+        message: CallbackMessage,
+    ) {
         let state = state_rc.clone();
         let mut state = state.borrow_mut();
         match message {

firmware/acid-firmware/src/util.rs

@@ -51,6 +51,7 @@ pub const fn get_file_name(path: &str) -> &str {
     }
 }
 
+#[allow(unused)]
 pub trait MutexExt<M, T> {
     type Guard<'a>
     where