/*
  Simd-types for parallel dsp processing.
  Aligned memory allocation for simd vectors.
 */

#pragma once
#include <cassert>
#include <cstdlib>

#if _MSC_VER
#define __finl __forceinline
#define __vecc __vectorcall
#else
#define __finl inline __attribute__((always_inline))
#define __vecc
#endif

#if defined(__SSE2__) || (defined(_M_AMD64) || defined(_M_X64)) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
#include "SimdTypes_sse2.h"
#elif defined(__arm64__) || defined(__aarch64__) || defined(_M_ARM64)
#include "SimdTypes_neon.h"
#else
#include "SimdTypes_scalar.h"
#endif

namespace staffpad::audio::simd {

/// reserve aligned memory. Needs to be freed with aligned_free()
inline void *aligned_malloc(size_t required_bytes, size_t alignment)
{
  auto offset = alignment - 1 + sizeof(void *);
  auto p1 = std::malloc(required_bytes + offset);
  if (p1 == nullptr)
    return nullptr;
  // figure out aligned position
  void *p2 = (void *)(((size_t)(p1) + offset) & ~(alignment - 1));
  // write malloced pointer in front of aligned data
  ((void **)p2)[-1] = p1;
  return p2;
}

/// free memory allocated with aligned_malloc
inline void aligned_free(void *p)
{
  if (p)
    free(((void **)p)[-1]);
}

/// create a c++ class at an memory-aligned spot that needs to be deleted using aligned_delete
template <typename cls>
inline cls *aligned_new(int alignment)
{
  void *mem = aligned_malloc(sizeof(cls), alignment);
  return new (mem) cls();
}

/** delete objects created using aligned_new */
template <typename cls>
inline void aligned_delete(cls *obj)
{
  if (obj != nullptr)
  {
    obj->~cls();
    aligned_free((void *)obj);
  }
}

template <typename T>
inline bool is_aligned(T *obj, int alignment)
{
    return (((size_t)obj) & (alignment - 1)) == 0;
}

/** this template allows to write float SIMD code with the supported operators the following way:

float *a_vec;
const float *b_vec;
perform_parallel_simd_aligned(a_vec, b_vec, 512, [](auto &a, auto &b) {
auto t = a;
a = 3.f * a + 12.f * b;
b = 0.25f * a + 3.f * b;
});
*/

// two buffers read/write
template <typename fnc>
__finl void perform_parallel_simd_aligned(float *a, float *b, int n, const fnc &f)
{
  // fnc& f needs to be a lambda of type [](auto &a, auto &b){}.
  // the autos will be float_x4/float
  constexpr int N = 4;
  constexpr int byte_size = sizeof(float);

  assert(is_aligned(a, N * byte_size) && is_aligned(b, N * byte_size));

  for (int i = 0; i <= n - N; i += N)
  {
    auto x = float_x4_load_aligned(a + i);
    auto y = float_x4_load_aligned(b + i);
    f(x, y);
    store_aligned(x, a + i);
    store_aligned(y, b + i);
  }
  // deal with last partial packet
  for (int i = n & (~(N - 1)); i < n; ++i)
    f(a[i], b[i]);
}

/// template for applying math to one data buffer
template <typename fnc>
__finl void perform_parallel_simd_aligned(float *a, int n, const fnc &f)
{
  // fnc& f needs to be a lambda of type [](auto &a){}.
  constexpr int N = 4;
  constexpr int byte_size = sizeof(float);
  assert(is_aligned(a, N * byte_size));

  for (int i = 0; i <= n - N; i += N)
  {
    auto x = float_x4_load_aligned(a + i);
    f(x);
    store_aligned(x, a + i);
  }
  // deal with last partial packet
  for (int i = n & (~(N - 1)); i < n; ++i)
    f(a[i]);
}

} // namespace staffpad::audio::simd
