1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use crate::pool::{Pool, PoolAllocator};
use aligned_buffer::UniqueAlignedBuffer;
use std::sync::Arc;

/// A policy for retaining buffers in the pool.
pub trait BufferRetentionPolicy<const ALIGNMENT: usize>: Copy {
	fn should_retain(&self, buffer: &UniqueAlignedBuffer<ALIGNMENT>) -> bool;
}

/// A policy that retains all buffers.
#[derive(Default, Clone, Copy)]
pub struct RetainAllRetentionPolicy;

impl<const ALIGNMENT: usize> BufferRetentionPolicy<ALIGNMENT> for RetainAllRetentionPolicy {
	#[inline(always)]
	fn should_retain(&self, _: &UniqueAlignedBuffer<ALIGNMENT>) -> bool {
		true
	}
}

/// A policy that retains buffers up to a maximum size.
#[derive(Default, Clone, Copy)]
pub struct ConstMaxSizeRetentionPolicy<const SIZE: usize>;

impl<const SIZE: usize, const ALIGNMENT: usize> BufferRetentionPolicy<ALIGNMENT>
	for ConstMaxSizeRetentionPolicy<SIZE>
{
	#[inline(always)]
	fn should_retain(&self, buffer: &UniqueAlignedBuffer<ALIGNMENT>) -> bool {
		buffer.len() <= SIZE
	}
}

#[derive(Default)]
struct BufferAllocator<P: BufferRetentionPolicy<ALIGNMENT>, const ALIGNMENT: usize> {
	policy: P,
}

impl<P: BufferRetentionPolicy<ALIGNMENT>, const ALIGNMENT: usize> BufferAllocator<P, ALIGNMENT> {
	fn new(policy: P) -> Self {
		Self { policy }
	}
}

impl<P: BufferRetentionPolicy<ALIGNMENT>, const ALIGNMENT: usize>
	PoolAllocator<UniqueAlignedBuffer<ALIGNMENT>> for BufferAllocator<P, ALIGNMENT>
{
	fn allocate(&self) -> UniqueAlignedBuffer<ALIGNMENT> {
		UniqueAlignedBuffer::new()
	}

	#[inline(always)]
	fn reset(&self, buf: &mut UniqueAlignedBuffer<ALIGNMENT>) {
		buf.clear();
	}

	#[inline(always)]
	fn is_valid(&self, buf: &UniqueAlignedBuffer<ALIGNMENT>) -> bool {
		self.policy.should_retain(buf)
	}
}

pub(crate) struct AlignedBufferPoolInner<
	P: BufferRetentionPolicy<ALIGNMENT>,
	const ALIGNMENT: usize,
> {
	pool: Pool<BufferAllocator<P, ALIGNMENT>, UniqueAlignedBuffer<ALIGNMENT>>,
}

impl<P: BufferRetentionPolicy<ALIGNMENT>, const ALIGNMENT: usize>
	AlignedBufferPoolInner<P, ALIGNMENT>
{
	pub fn new(policy: P, capacity: usize) -> Self {
		Self {
			pool: Pool::new(BufferAllocator::new(policy), capacity),
		}
	}

	/// Gets a buffer from the pool. If the pool is empty, a new buffer is
	/// allocated and returned.
	#[inline]
	pub fn get(&self) -> UniqueAlignedBuffer<ALIGNMENT> {
		self.pool.get()
	}

	/// Recycles a buffer, putting it back into the pool.
	#[inline]
	pub fn recycle(&self, buffer: UniqueAlignedBuffer<ALIGNMENT>) {
		self.pool.recycle(buffer)
	}
}

/// A pool for allocating and recycling aligned buffers.
pub struct AlignedBufferPool<P: BufferRetentionPolicy<ALIGNMENT>, const ALIGNMENT: usize> {
	inner: Arc<AlignedBufferPoolInner<P, ALIGNMENT>>,
}

impl<P: BufferRetentionPolicy<ALIGNMENT>, const ALIGNMENT: usize> AlignedBufferPool<P, ALIGNMENT> {
	pub fn new(policy: P, capacity: usize) -> Self {
		Self {
			inner: Arc::new(AlignedBufferPoolInner::new(policy, capacity)),
		}
	}

	/// Gets a buffer from the pool. If the pool is empty, a new buffer is
	/// allocated and returned.
	#[inline]
	pub fn get(&self) -> UniqueAlignedBuffer<ALIGNMENT> {
		self.inner.get()
	}

	/// Recycles a buffer, putting it back into the pool.
	#[inline]
	pub fn recycle(&self, buffer: UniqueAlignedBuffer<ALIGNMENT>) {
		self.inner.recycle(buffer)
	}
}

impl<P: BufferRetentionPolicy<ALIGNMENT>, const ALIGNMENT: usize> AlignedBufferPool<P, ALIGNMENT>
where
	P: Default,
{
	pub fn with_capacity(capacity: usize) -> Self {
		Self::new(P::default(), capacity)
	}
}