From 226901861adbc0def1e4a756c6169b6fc0c6d88c Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Thu, 29 Jul 2021 23:42:27 +0100
Subject: [PATCH] Start adding stack-allocation helper functions. (ext.rs)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fortune for rsh's current commit: Half curse − 半凶
---
 Cargo.lock  | 11 +++++++++++
 Cargo.toml  |  1 +
 src/ext.rs  | 38 ++++++++++++++++++++++++++++++++++++++
 src/main.rs |  1 +
 4 files changed, 51 insertions(+)
 create mode 100644 src/ext.rs

diff --git a/Cargo.lock b/Cargo.lock
index a558574..5904eca 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -780,6 +780,7 @@ dependencies = [
  "pin-project",
  "serde",
  "serde_cbor",
+ "stackalloc",
  "tokio 0.2.25",
  "tokio-uring",
  "uuid",
@@ -894,6 +895,16 @@ version = "1.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fe0f37c9e8f3c5a4a66ad655a93c74daac4ad00c441533bf5c6e7990bb42604e"
 
+[[package]]
+name = "stackalloc"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4f5c9dd3feb8a4adc8eae861e5f48862a92f9a9f38cf8fc99b92fc6ec016121"
+dependencies = [
+ "cc",
+ "rustc_version",
+]
+
 [[package]]
 name = "subtle"
 version = "2.4.1"
diff --git a/Cargo.toml b/Cargo.toml
index 5a0a2a9..f845cd9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,6 +12,7 @@ cryptohelpers = { version = "1.8", features = ["serialise", "full"] }
 pin-project = "1.0.8"
 serde = { version = "1.0.126", features = ["derive"] }
 serde_cbor = "0.11.1"
+stackalloc = "1.1.1"
 tokio = { version = "0.2", features = ["full"] }
 tokio-uring = "0.1.0"
 uuid = { version = "0.8.2", features = ["v4", "serde"] }
diff --git a/src/ext.rs b/src/ext.rs
new file mode 100644
index 0000000..c73552a
--- /dev/null
+++ b/src/ext.rs
@@ -0,0 +1,38 @@
+//! Extensions
+use super::*;
+use std::mem::{self, MaybeUninit};
+use std::iter;
+
+/// Max size of memory allowed to be allocated on the stack.
+pub const STACK_MEM_ALLOC_MAX: usize = 4096;
+
+pub fn vec_uninit<T>(sz: usize) -> Vec<MaybeUninit<T>>
+{
+    let mut mem: Vec<T> = Vec::with_capacity(sz);
+    unsafe {
+	mem.set_len(sz);
+	mem::transmute(mem)
+    }
+}
+
+/// Allocate a local buffer initialised with `init`.
+pub fn alloc_local_with<T, U>(sz: usize, init: impl FnMut() -> T, within: impl FnOnce(&mut [T]) -> U) -> U
+{
+    if sz > STACK_MEM_ALLOC_MAX {
+	let mut memory: Vec<T> = iter::repeat_with(init).take(sz).collect();
+	within(&mut memory[..])
+    } else {
+	stackalloc::stackalloc_with(sz, init, within)
+    }
+}
+
+/// Allocate a local buffer initialised with `init`.
+pub fn alloc_local<T: Clone, U>(sz: usize, init: T, within: impl FnOnce(&mut [T]) -> U) -> U
+{
+    if sz > STACK_MEM_ALLOC_MAX {
+	let mut memory: Vec<T> = iter::repeat(init).take(sz).collect();
+	within(&mut memory[..])
+    } else {
+	stackalloc::stackalloc(sz, init, within)
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index ee8dcf5..48135c4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,6 +19,7 @@ use std::convert::{
     TryInto,
 };
 
+mod ext; use ext::*;
 mod message;
 
 #[tokio::main]