1 #pragma once
2 
3 #ifndef ZIMG_ALLOC_H_
4 #define ZIMG_ALLOC_H_
5 
6 #include <cstddef>
7 #include <stdexcept>
8 #include <vector>
9 #include "align.h"
10 #include "checked_int.h"
11 #include "except.h"
12 
13 #ifdef _WIN32
14   #include <malloc.h>
zimg_x_aligned_malloc(size_t size,size_t alignment)15   inline void *zimg_x_aligned_malloc(size_t size, size_t alignment) { return _aligned_malloc(size, alignment); }
zimg_x_aligned_free(void * ptr)16   inline void zimg_x_aligned_free(void *ptr) { _aligned_free(ptr); }
17 #else
18   #include <stdlib.h>
zimg_x_aligned_malloc(size_t size,size_t alignment)19   inline void *zimg_x_aligned_malloc(size_t size, size_t alignment) { void *p; if (posix_memalign(&p, alignment, size)) return nullptr; else return p; }
zimg_x_aligned_free(void * ptr)20   inline void zimg_x_aligned_free(void *ptr) { free(ptr); }
21 #endif
22 
23 namespace zimg {
24 
25 /**
26  * Simple allocator that increments a base pointer.
27  * This allocator is not STL compliant.
28  */
29 class LinearAllocator {
30 	char *m_ptr;
31 	size_t m_count;
32 
33 	template <class T>
increment_and_return(size_t n)34 	T *increment_and_return(size_t n)
35 	{
36 		char *ptr = m_ptr;
37 		m_ptr += n;
38 		m_count += n;
39 		return reinterpret_cast<T *>(ptr);
40 	}
41 public:
42 	/**
43 	 * Initialize a LinearAllocator with a given buffer.
44 	 *
45 	 * @param ptr pointer to buffer
46 	 */
LinearAllocator(void * ptr)47 	LinearAllocator(void *ptr) noexcept :
48 		m_ptr{ static_cast<char *>(ptr) },
49 		m_count{}
50 	{}
51 
52 	/**
53 	 * Allocate a given number of bytes.
54 	 *
55 	 * @tparam T buffer type
56 	 * @param bytes number of bytes
57 	 * @return pointer to buffer
58 	 */
59 	template <class T = void>
allocate(size_t bytes)60 	T *allocate(size_t bytes)
61 	{
62 		return increment_and_return<T>(ceil_n(bytes, ALIGNMENT));
63 	}
64 
65 	/**
66 	 * Allocate a given number of objects.
67 	 *
68 	 * @tparam T buffer type
69 	 * @param count number of objects
70 	 * @return pointer to buffer
71 	 */
72 	template <class T>
allocate_n(size_t count)73 	T *allocate_n(size_t count)
74 	{
75 		return allocate<T>(count * sizeof(T));
76 	}
77 
78 	/**
79 	 * Get the number of bytes allocated.
80 	 *
81 	 * @return allocated size
82 	 */
count()83 	size_t count() const noexcept
84 	{
85 		return m_count;
86 	}
87 };
88 
89 /**
90  * Fake allocator that tracks the amount allocated.
91  * The pointers returned must not be dereferenced or incremented.
92  */
93 class FakeAllocator {
94 	checked_size_t m_count;
95 public:
96 	/**
97 	 * Initialize a FakeAllocator.
98 	 */
FakeAllocator()99 	FakeAllocator() noexcept : m_count{} {}
100 
101 	/**
102 	 * Allocate a given number of bytes.
103 	 *
104 	 * @tparam T buffer type
105 	 * @param bytes number of bytes
106 	 * @return nullptr
107 	 * @throw error::OutOfMemory if total allocation exceeds SIZE_MAX
108 	 */
109 	template <class T = void>
allocate(checked_size_t bytes)110 	T *allocate(checked_size_t bytes)
111 	{
112 		try {
113 			m_count += ceil_n(bytes, ALIGNMENT);
114 		} catch (const std::overflow_error &) {
115 			error::throw_<error::OutOfMemory>();
116 		}
117 
118 		return nullptr;
119 	}
120 
121 	/**
122 	 * Allocate a given number of objects.
123 	 *
124 	 * @tparam T buffer type
125 	 * @param count number of objects
126 	 * @return nullptr
127 	 * @throw error::OutOfMemory if total allocation exceeds SIZE_MAX
128 	 */
129 	template <class T>
allocate_n(checked_size_t count)130 	T *allocate_n(checked_size_t count)
131 	{
132 		return allocate<T>(count * sizeof(T));
133 	}
134 
135 	/**
136 	 * Get the number of bytes allocated.
137 	 *
138 	 * @return allocated size
139 	 */
count()140 	size_t count() const noexcept
141 	{
142 		return m_count.get();
143 	}
144 };
145 
146 /**
147  * STL allocator class which returns aligned buffers.
148  *
149  * @tparam T type of object to allocate
150  */
151 template <class T>
152 struct AlignedAllocator {
153 	typedef T value_type;
154 
155 	AlignedAllocator() = default;
156 
157 	template <class U>
AlignedAllocatorAlignedAllocator158 	AlignedAllocator(const AlignedAllocator<U> &) noexcept {}
159 
allocateAlignedAllocator160 	T *allocate(size_t n) const
161 	{
162 		T *ptr = static_cast<T *>(zimg_x_aligned_malloc(n * sizeof(T), ALIGNMENT));
163 
164 		if (!ptr)
165 			throw std::bad_alloc{};
166 
167 		return ptr;
168 	}
169 
deallocateAlignedAllocator170 	void deallocate(void *ptr, size_t) const noexcept
171 	{
172 		zimg_x_aligned_free(ptr);
173 	}
174 
175 	bool operator==(const AlignedAllocator &) const noexcept { return true; }
176 	bool operator!=(const AlignedAllocator &) const noexcept { return false; }
177 };
178 
179 /**
180  * std::vector specialization using AlignedAllocator.
181  */
182 template <class T>
183 using AlignedVector = std::vector<T, AlignedAllocator<T>>;
184 
185 } // namespace zimg
186 
187 #endif // ZIMG_ALLOC_H_
188