From bcef8ed7a92bbf0d5069349e673438634aab2de7 Mon Sep 17 00:00:00 2001
From: Luca Versari <veluca93@gmail.com>
Date: Wed, 19 Nov 2025 23:24:58 +0100
Subject: [PATCH] Use talc to handle allocations.

---
 pixie-uefi/Cargo.lock          | 35 ++++++++++++++++++++++++++++++++
 pixie-uefi/Cargo.toml          |  4 +++-
 pixie-uefi/src/os/allocator.rs | 37 ++++++++++++++++++++++++++++++++++
 pixie-uefi/src/os/mod.rs       |  1 +
 4 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 pixie-uefi/src/os/allocator.rs

diff --git a/pixie-uefi/Cargo.lock b/pixie-uefi/Cargo.lock
index f6f6ab88..ea4f7c35 100644
--- a/pixie-uefi/Cargo.lock
+++ b/pixie-uefi/Cargo.lock
@@ -306,6 +306,15 @@ dependencies = [
  "stable_deref_trait",
 ]
 
+[[package]]
+name = "lock_api"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
+dependencies = [
+ "scopeguard",
+]
+
 [[package]]
 name = "log"
 version = "0.4.28"
@@ -374,6 +383,8 @@ dependencies = [
  "rand",
  "rand_xoshiro",
  "smoltcp",
+ "spin",
+ "talc",
  "uefi",
 ]
 
@@ -482,6 +493,12 @@ dependencies = [
  "winapi-util",
 ]
 
+[[package]]
+name = "scopeguard"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
+
 [[package]]
 name = "serde"
 version = "1.0.228"
@@ -532,6 +549,15 @@ dependencies = [
  "managed",
 ]
 
+[[package]]
+name = "spin"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d5fe4ccb98d9c292d56fec89a5e07da7fc4cf0dc11e156b41793132775d3e591"
+dependencies = [
+ "lock_api",
+]
+
 [[package]]
 name = "stable_deref_trait"
 version = "1.2.0"
@@ -560,6 +586,15 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "talc"
+version = "4.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a3ae828aa394de34c7de08f522d1b86bd1c182c668d27da69caadda00590f26d"
+dependencies = [
+ "lock_api",
+]
+
 [[package]]
 name = "thiserror"
 version = "2.0.17"
diff --git a/pixie-uefi/Cargo.toml b/pixie-uefi/Cargo.toml
index 6cae3362..2cb71a3f 100644
--- a/pixie-uefi/Cargo.toml
+++ b/pixie-uefi/Cargo.toml
@@ -27,7 +27,9 @@ postcard = { version = "1.1.3", default-features = false, features = ["alloc"] }
 rand = { version = "0.8.5", default-features = false }
 rand_xoshiro = { version = "0.6.0", default-features = false }
 smoltcp = { version = "0.12.0", default-features = false, features = ["alloc", "proto-ipv4", "medium-ethernet", "socket-udp", "socket-tcp", "socket-dhcpv4", "async", "socket-tcp-cubic"] }
-uefi = { version = "0.33.0", features = ["alloc", "global_allocator", "panic_handler"] }
+spin = "0.10.0"
+talc = { version = "4.4.3", default-features = false, features = ["lock_api"] }
+uefi = { version = "0.33.0", features = ["panic_handler", "alloc"] }
 
 [dependencies.pixie-shared]
 path = "../pixie-shared"
diff --git a/pixie-uefi/src/os/allocator.rs b/pixie-uefi/src/os/allocator.rs
new file mode 100644
index 00000000..30d0a618
--- /dev/null
+++ b/pixie-uefi/src/os/allocator.rs
@@ -0,0 +1,37 @@
+use talc::*;
+use uefi::boot::MemoryType;
+
+const MAX_BLOCK_SIZE: usize = 1 << 30;
+
+struct AllocOnOom {
+    next_block_size: usize,
+}
+
+impl OomHandler for AllocOnOom {
+    fn handle_oom(talc: &mut Talc<Self>, layout: core::alloc::Layout) -> Result<(), ()> {
+        let bs = talc.oom_handler.next_block_size;
+        talc.oom_handler.next_block_size =
+            (talc.oom_handler.next_block_size * 2).min(MAX_BLOCK_SIZE);
+        match uefi::boot::allocate_pool(MemoryType::LOADER_DATA, bs) {
+            Err(e) => {
+                uefi::println!(
+                    "Cannot allocate new block ({e}). Triggered by allocation {layout:?}"
+                );
+                Err(())
+            }
+            Ok(ptr) => {
+                let span = talc::Span::from_base_size(ptr.as_ptr(), bs);
+                // SAFETY: the memory was just allocated, so we have exclusive access to it
+                // and we can transfer ownership.
+                unsafe { talc.claim(span) }?;
+                Ok(())
+            }
+        }
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: Talck<spin::Mutex<()>, AllocOnOom> = Talc::new(AllocOnOom {
+    next_block_size: 1 << 20,
+})
+.lock();
diff --git a/pixie-uefi/src/os/mod.rs b/pixie-uefi/src/os/mod.rs
index 27af756e..86f24a7d 100644
--- a/pixie-uefi/src/os/mod.rs
+++ b/pixie-uefi/src/os/mod.rs
@@ -44,6 +44,7 @@ use uefi::{
     CStr16, Event, Handle, Status,
 };
 
+mod allocator;
 mod boot_options;
 pub mod disk;
 pub mod error;