/*
 * Copyright (c) 2023 by FlashInfer team.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef FLASHINFER_ALLOCATOR_H_
#define FLASHINFER_ALLOCATOR_H_

#include <limits>
#include <memory>
#include <sstream>

#include "exception.h"

namespace flashinfer {

// create a function that returns T* from base pointer and offset
template <typename T>
T* GetPtrFromBaseOffset(void* base_ptr, int64_t offset) {
  return reinterpret_cast<T*>(reinterpret_cast<char*>(base_ptr) + offset);
}

struct AlignedAllocator {
  void* base_ptr;
  void* cur_ptr;
  size_t remaining_space;
  size_t allocated_space;
  bool count_only;
  AlignedAllocator(void* buf, size_t space)
      : base_ptr(buf),
        cur_ptr(buf),
        remaining_space(space),
        allocated_space(0),
        count_only(false) {}
  // Counting mode starts from offset 0 and assumes caller-owned workspace
  // buffers are aligned to every allocation alignment requested below.
  explicit AlignedAllocator(size_t space = std::numeric_limits<size_t>::max())
      : base_ptr(nullptr),
        cur_ptr(nullptr),
        remaining_space(space),
        allocated_space(0),
        count_only(true) {}
  template <typename T>
  T* aligned_alloc(size_t size, size_t alignment, std::string name) {
    if (count_only) {
      std::ostringstream oss;
      oss << "AlignedAllocator::aligned_alloc cannot return a pointer in counting mode for "
          << name;
      FLASHINFER_ERROR(oss.str());
    }
    if (std::align(alignment, size, cur_ptr, remaining_space)) {
      T* result = reinterpret_cast<T*>(cur_ptr);
      cur_ptr = (char*)cur_ptr + size;
      remaining_space -= size;
      return result;
    } else {
      std::ostringstream oss;
      oss << "Buffer overflow when allocating memory for " << name << " with size " << size
          << " and alignment " << alignment << ", but only " << remaining_space
          << " bytes available in AlignedAllocator. Increase the workspace buffer size.";
      FLASHINFER_ERROR(oss.str());
    }
    return nullptr;
  }

  size_t aligned_alloc_offset(size_t size, size_t alignment, std::string name) {
    if (count_only) {
      size_t padding = 0;
      if (alignment > 1) {
        padding = (alignment - (allocated_space % alignment)) % alignment;
      }
      if (padding > remaining_space || size > remaining_space - padding) {
        std::ostringstream oss;
        oss << "Buffer overflow when allocating memory for " << name << " with size " << size
            << " and alignment " << alignment << ", but only " << remaining_space
            << " bytes available in AlignedAllocator. Increase the workspace buffer size.";
        FLASHINFER_ERROR(oss.str());
      }
      const size_t result = allocated_space + padding;
      allocated_space = result + size;
      remaining_space -= padding + size;
      return result;
    }
    return (char*)aligned_alloc<char>(size, alignment, name) - (char*)base_ptr;
  }

  size_t num_allocated_bytes() {
    return count_only ? allocated_space : (char*)cur_ptr - (char*)base_ptr;
  }
};

}  // namespace flashinfer

#endif  // FLASHINFER_ALLOCATOR_H_
