Alicho/src/simd/aligned_allocator.h

#pragma once
#include <memory>
#include <cstdlib>
#include <new>
#include <cstddef>
#include <stdexcept> // ASSERT_ALIGNED宏需要

constexpr size_t ALIGNMENT_SSE    = 16; // SSE要求16字节对齐
constexpr size_t ALIGNMENT_AVX    = 32; // AVX要求32字节对齐
constexpr size_t ALIGNMENT_AVX512 = 64; // AVX-512要求64字节对齐
constexpr size_t ALIGNMENT_CACHE  = 64; // CPU缓存行对齐(通常为64字节)

inline auto aligned_malloc(size_t size, size_t alignment) -> void* {
	if (alignment == 0 || (alignment & (alignment - 1)) != 0) {
		// 对齐值必须是2的幂
		return nullptr;
	}

#if ALICHO_PLATFORM_WINDOWS
	return _aligned_malloc(size, alignment);
#elif ALICHO_PLATFORM_POSIX || ALICHO_PLATFORM_UNIX
	void* ptr = nullptr;
	if (posix_memalign(&ptr, alignment, size) != 0) {
		return nullptr;
	}
	return ptr;
#else
	// 回退实现：手动对齐
	// 分配额外空间来存储原始指针和进行对齐
	size_t total_size = size + alignment + sizeof(void*);
	void* raw_ptr = std::malloc(total_size);
	if (!raw_ptr) {
		return nullptr;
	}

	// 计算对齐后的地址
	uintptr_t raw_addr = reinterpret_cast<uintptr_t>(raw_ptr);
	uintptr_t aligned_addr = (raw_addr + sizeof(void*) + alignment - 1) & ~(alignment - 1);
	void* aligned_ptr = reinterpret_cast<void*>(aligned_addr);

	// 在对齐地址前存储原始指针
	(reinterpret_cast<void**>(aligned_ptr))[-1] = raw_ptr;

	return aligned_ptr;
#endif
}

inline void aligned_free(void* ptr) {
	if (!ptr) {
		return;
	}
#if ALICHO_PLATFORM_WINDOWS
	_aligned_free(ptr);
#elif ALICHO_PLATFORM_POSIX || ALICHO_PLATFORM_UNIX
	std::free(ptr);
#else
	// 回退实现：获取原始指针并释放
	void* raw_ptr = (reinterpret_cast<void**>(ptr))[-1];
	std::free(raw_ptr);
#endif
}

// 对齐分配器模板类
template<typename T, size_t Alignment = ALIGNMENT_AVX>
class aligned_allocator {
public:
	using value_type = T;
	using pointer = T*;
	using const_pointer = const T*;
	using reference = T&;
	using const_reference = const T&;
	using size_type = std::size_t;
	using difference_type = std::ptrdiff_t;

	template<typename U>
	struct rebind {
		using other = aligned_allocator<U, Alignment>;
	};

	aligned_allocator() noexcept = default;

	template<typename U>
	aligned_allocator(const aligned_allocator<U, Alignment>&) noexcept {}

	auto allocate(size_type n) -> pointer {
		if (n == 0) return nullptr;

		size_type size = n * sizeof(T);
		void* ptr = aligned_malloc(size, Alignment);

		if (!ptr) {
			throw std::bad_alloc();
		}

		return static_cast<pointer>(ptr);
	}

	void deallocate(pointer p, size_type) noexcept {
		aligned_free(p);
	}

	template<typename U, typename... Args>
	void construct(U* p, Args&&... args) {
		::new(static_cast<void*>(p)) U(std::forward<Args>(args)...);
	}

	template<typename U>
	void destroy(U* p) {
		p->~u();
	}

	static auto max_size() noexcept {
		return std::numeric_limits<size_type>::max() / sizeof(T);
	}
};

template<typename T1, size_t A1, typename T2, size_t A2>
bool operator==(const aligned_allocator<T1, A1>&, const aligned_allocator<T2, A2>&) noexcept {
	return A1 == A2;
}

template<typename T1, size_t A1, typename T2, size_t A2>
bool operator!=(const aligned_allocator<T1, A1>&, const aligned_allocator<T2, A2>&) noexcept {
	return A1 != A2;
}

// 类型别名，方便使用不同对齐方式的分配器
template<typename T>
using sse_aligned_allocator = aligned_allocator<T, ALIGNMENT_SSE>;
template<typename T>
using avx_aligned_allocator = aligned_allocator<T, ALIGNMENT_AVX>;
template<typename T>
using avx512_aligned_allocator = aligned_allocator<T, ALIGNMENT_AVX512>;
template<typename T>
using cache_aligned_allocator = aligned_allocator<T, ALIGNMENT_CACHE>;

template<size_t Alignment>
auto is_aligned(void* ptr) -> bool {
	return (reinterpret_cast<uintptr_t>(ptr) % Alignment) == 0;
}

inline auto is_aligned(const void* ptr, size_t alignment) -> bool {
	return (reinterpret_cast<uintptr_t>(ptr) % alignment) == 0;
}

inline auto align_size(size_t size, size_t alignment) -> size_t {
	return (size + alignment - 1) & ~(alignment - 1);
}

template<size_t Alignment>
auto align_pointer(void* ptr) -> void* {
	const auto addr = reinterpret_cast<uintptr_t>(ptr);
	const auto aligned_addr = (addr + Alignment - 1) & ~(Alignment - 1);
	return reinterpret_cast<void*>(aligned_addr);
}

template<typename T, size_t Alignment>
class aligned_buffer {
public:
	aligned_buffer() = default;

	explicit aligned_buffer(size_t size) : size_(size) {
		allocate(size);
	}

	void allocate(size_t new_size) {
		if (data_) {
			deallocate();
		}
		
		if (new_size == 0) {
			data_ = nullptr;
			size_ = 0;
			return;
		}
		
		data_ = static_cast<T*>(aligned_malloc(new_size * sizeof(T), Alignment));
		if (!data_) {
			throw std::bad_alloc();
		}
		
		size_ = new_size;
		
		// 对于非POD类型，需要构造对象
		if constexpr (!std::is_trivially_constructible_v<T>) {
			for (size_t i = 0; i < size_; ++i) {
				new(&data_[i]) T();
			}
		}
	}

	void deallocate() {
		if (!data_)
			return;
		// 对于非POD类型，需要析构对象
		if constexpr (!std::is_trivially_destructible_v<T>) {
			for (size_t i = 0; i < size_; ++i) {
				data_[i].~T();
			}
		}
		aligned_free(data_);
		data_ = nullptr;
		size_ = 0;
	}

	void resize(size_t size) {
		if (size == size_)
			return;
		allocate(size);
	}

	auto data() noexcept { return data_; }
	auto data() const noexcept { return data_; }
	auto size() const noexcept { return size_; }
	auto empty() const noexcept { return size_ == 0; }

	auto& operator[](size_t index) noexcept { return data_[index]; }
	const auto& operator[](size_t index) const noexcept { return data_[index]; }

	auto begin() noexcept { return data_; }
	auto end() noexcept { return data_ + size_; }
	auto begin() const noexcept { return data_; }
	auto end() const noexcept { return data_ + size_; }

	[[nodiscard]] auto is_properly_aligned() const noexcept -> bool {
		return is_aligned<Alignment>(data_);
	}
private:
	T* data_ = nullptr;
	size_t size_ = 0;
};

#if ALICHO_DEBUG
#define ASSERT_ALIGNED(ptr, alignment) \
	do { \
		if (!is_aligned(ptr, alignment)) { \
			throw std::runtime_error("Pointer is not properly aligned"); \
		} \
	} while (0)
#else
#define ASSERT_ALIGNED(ptr, alignment)
#endif