From a97737a869486cd308c6d542e5f2d89a7b7d1538 Mon Sep 17 00:00:00 2001
From: Dylan Plecki <dylan@plecki.net>
Date: Wed, 17 Jan 2024 17:46:00 -0500
Subject: [PATCH] Add optional feature to support the nightly/unstable
 allocator API

---
 Cargo.toml           |   1 +
 src/buf/buf_impl.rs  |   6 +++
 src/buf/buf_mut.rs   | 120 +++++++++++++++++++++++--------------------
 src/buf/macros.rs    |  22 ++++++++
 src/buf/mod.rs       |   1 +
 src/buf/vec_deque.rs |  29 ++++++-----
 src/lib.rs           |   1 +
 7 files changed, 111 insertions(+), 69 deletions(-)
 create mode 100644 src/buf/macros.rs

diff --git a/Cargo.toml b/Cargo.toml
index 127d81dd5..2cabd8eb7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,6 +21,7 @@ categories = ["network-programming", "data-structures"]
 [features]
 default = ["std"]
 std = []
+allocator_api = []
 
 [dependencies]
 serde = { version = "1.0.60", optional = true, default-features = false, features = ["alloc"] }
diff --git a/src/buf/buf_impl.rs b/src/buf/buf_impl.rs
index 9367eb2df..1a1566c0f 100644
--- a/src/buf/buf_impl.rs
+++ b/src/buf/buf_impl.rs
@@ -1389,10 +1389,16 @@ impl<T: Buf + ?Sized> Buf for &mut T {
     deref_forward_buf!();
 }
 
+#[cfg(not(feature = "allocator_api"))]
 impl<T: Buf + ?Sized> Buf for Box<T> {
     deref_forward_buf!();
 }
 
+#[cfg(feature = "allocator_api")]
+impl<T: Buf + ?Sized, A: core::alloc::Allocator> Buf for Box<T, A> {
+    deref_forward_buf!();
+}
+
 impl Buf for &[u8] {
     #[inline]
     fn remaining(&self) -> usize {
diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs
index 304e11b13..bab9381d3 100644
--- a/src/buf/buf_mut.rs
+++ b/src/buf/buf_mut.rs
@@ -1,7 +1,7 @@
 use crate::buf::{limit, Chain, Limit, UninitSlice};
 #[cfg(feature = "std")]
 use crate::buf::{writer, Writer};
-use crate::{panic_advance, panic_does_not_fit};
+use crate::{impl_with_allocator, panic_advance, panic_does_not_fit};
 
 use core::{mem, ptr, usize};
 
@@ -1469,10 +1469,16 @@ unsafe impl<T: BufMut + ?Sized> BufMut for &mut T {
     deref_forward_bufmut!();
 }
 
+#[cfg(not(feature = "allocator_api"))]
 unsafe impl<T: BufMut + ?Sized> BufMut for Box<T> {
     deref_forward_bufmut!();
 }
 
+#[cfg(feature = "allocator_api")]
+unsafe impl<T: BufMut + ?Sized, A: core::alloc::Allocator> BufMut for Box<T, A> {
+    deref_forward_bufmut!();
+}
+
 unsafe impl BufMut for &mut [u8] {
     #[inline]
     fn remaining_mut(&self) -> usize {
@@ -1569,70 +1575,72 @@ unsafe impl BufMut for &mut [core::mem::MaybeUninit<u8>] {
     }
 }
 
-unsafe impl BufMut for Vec<u8> {
-    #[inline]
-    fn remaining_mut(&self) -> usize {
-        // A vector can never have more than isize::MAX bytes
-        core::isize::MAX as usize - self.len()
-    }
-
-    #[inline]
-    unsafe fn advance_mut(&mut self, cnt: usize) {
-        let len = self.len();
-        let remaining = self.capacity() - len;
-
-        if remaining < cnt {
-            panic_advance(cnt, remaining);
+impl_with_allocator! {
+    unsafe impl BufMut for Vec<u8> {
+        #[inline]
+        fn remaining_mut(&self) -> usize {
+            // A vector can never have more than isize::MAX bytes
+            core::isize::MAX as usize - self.len()
         }
 
-        // Addition will not overflow since the sum is at most the capacity.
-        self.set_len(len + cnt);
-    }
-
-    #[inline]
-    fn chunk_mut(&mut self) -> &mut UninitSlice {
-        if self.capacity() == self.len() {
-            self.reserve(64); // Grow the vec
-        }
+        #[inline]
+        unsafe fn advance_mut(&mut self, cnt: usize) {
+            let len = self.len();
+            let remaining = self.capacity() - len;
 
-        let cap = self.capacity();
-        let len = self.len();
+            if remaining < cnt {
+                panic_advance(cnt, remaining);
+            }
 
-        let ptr = self.as_mut_ptr();
-        // SAFETY: Since `ptr` is valid for `cap` bytes, `ptr.add(len)` must be
-        // valid for `cap - len` bytes. The subtraction will not underflow since
-        // `len <= cap`.
-        unsafe { UninitSlice::from_raw_parts_mut(ptr.add(len), cap - len) }
-    }
+            // Addition will not overflow since the sum is at most the capacity.
+            self.set_len(len + cnt);
+        }
 
-    // Specialize these methods so they can skip checking `remaining_mut`
-    // and `advance_mut`.
-    #[inline]
-    fn put<T: super::Buf>(&mut self, mut src: T)
-    where
-        Self: Sized,
-    {
-        // In case the src isn't contiguous, reserve upfront.
-        self.reserve(src.remaining());
+        #[inline]
+        fn chunk_mut(&mut self) -> &mut UninitSlice {
+            if self.capacity() == self.len() {
+                self.reserve(64); // Grow the vec
+            }
+
+            let cap = self.capacity();
+            let len = self.len();
+
+            let ptr = self.as_mut_ptr();
+            // SAFETY: Since `ptr` is valid for `cap` bytes, `ptr.add(len)` must be
+            // valid for `cap - len` bytes. The subtraction will not underflow since
+            // `len <= cap`.
+            unsafe { UninitSlice::from_raw_parts_mut(ptr.add(len), cap - len) }
+        }
 
-        while src.has_remaining() {
-            let s = src.chunk();
-            let l = s.len();
-            self.extend_from_slice(s);
-            src.advance(l);
+        // Specialize these methods so they can skip checking `remaining_mut`
+        // and `advance_mut`.
+        #[inline]
+        fn put<T: super::Buf>(&mut self, mut src: T)
+        where
+            Self: Sized,
+        {
+            // In case the src isn't contiguous, reserve upfront.
+            self.reserve(src.remaining());
+
+            while src.has_remaining() {
+                let s = src.chunk();
+                let l = s.len();
+                self.extend_from_slice(s);
+                src.advance(l);
+            }
         }
-    }
 
-    #[inline]
-    fn put_slice(&mut self, src: &[u8]) {
-        self.extend_from_slice(src);
-    }
+        #[inline]
+        fn put_slice(&mut self, src: &[u8]) {
+            self.extend_from_slice(src);
+        }
 
-    #[inline]
-    fn put_bytes(&mut self, val: u8, cnt: usize) {
-        // If the addition overflows, then the `resize` will fail.
-        let new_len = self.len().saturating_add(cnt);
-        self.resize(new_len, val);
+        #[inline]
+        fn put_bytes(&mut self, val: u8, cnt: usize) {
+            // If the addition overflows, then the `resize` will fail.
+            let new_len = self.len().saturating_add(cnt);
+            self.resize(new_len, val);
+        }
     }
 }
 
diff --git a/src/buf/macros.rs b/src/buf/macros.rs
new file mode 100644
index 000000000..2372198b1
--- /dev/null
+++ b/src/buf/macros.rs
@@ -0,0 +1,22 @@
+/// impl_with_allocator implements a specified trait for a concrete type that
+/// is expected to take some number of generic arguments and optionally a
+/// trailing allocator argument of type [core::alloc::Allocator] only if the
+/// unstable `allocator_api` feature is enabled.
+#[macro_export]
+macro_rules! impl_with_allocator {
+    { impl $interface:ident for $concrete:ident<$( $generic_args:ty ),*> $implementation:tt } => {
+        #[cfg(not(feature = "allocator_api"))]
+        impl $interface for $concrete<$( $generic_args ),*> $implementation
+
+        #[cfg(feature = "allocator_api")]
+        impl<A: core::alloc::Allocator> $interface for $concrete<$( $generic_args ),*, A> $implementation
+    };
+
+    { unsafe impl $interface:ident for $concrete:ident<$( $generic_args:ty ),*> $implementation:tt } => {
+        #[cfg(not(feature = "allocator_api"))]
+        unsafe impl $interface for $concrete<$( $generic_args ),*> $implementation
+
+        #[cfg(feature = "allocator_api")]
+        unsafe impl<A: core::alloc::Allocator> $interface for $concrete<$( $generic_args ),*, A> $implementation
+    };
+}
diff --git a/src/buf/mod.rs b/src/buf/mod.rs
index 1bf0a47e8..6593b5fd1 100644
--- a/src/buf/mod.rs
+++ b/src/buf/mod.rs
@@ -19,6 +19,7 @@ mod buf_mut;
 mod chain;
 mod iter;
 mod limit;
+mod macros;
 #[cfg(feature = "std")]
 mod reader;
 mod take;
diff --git a/src/buf/vec_deque.rs b/src/buf/vec_deque.rs
index 263167e83..071b96845 100644
--- a/src/buf/vec_deque.rs
+++ b/src/buf/vec_deque.rs
@@ -1,22 +1,25 @@
+use crate::impl_with_allocator;
 use alloc::collections::VecDeque;
 
 use super::Buf;
 
-impl Buf for VecDeque<u8> {
-    fn remaining(&self) -> usize {
-        self.len()
-    }
+impl_with_allocator! {
+    impl Buf for VecDeque<u8> {
+        fn remaining(&self) -> usize {
+            self.len()
+        }
 
-    fn chunk(&self) -> &[u8] {
-        let (s1, s2) = self.as_slices();
-        if s1.is_empty() {
-            s2
-        } else {
-            s1
+        fn chunk(&self) -> &[u8] {
+            let (s1, s2) = self.as_slices();
+            if s1.is_empty() {
+                s2
+            } else {
+                s1
+            }
         }
-    }
 
-    fn advance(&mut self, cnt: usize) {
-        self.drain(..cnt);
+        fn advance(&mut self, cnt: usize) {
+            self.drain(..cnt);
+        }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 1b3e6fc40..1923e8740 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,6 +4,7 @@
     attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_variables))
 ))]
 #![no_std]
+#![cfg_attr(feature = "allocator_api", feature(allocator_api))]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 
 //! Provides abstractions for working with bytes.