// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 IMPALA_RUNTIME_FREE_POOL_H
#define IMPALA_RUNTIME_FREE_POOL_H

#include <stdio.h>
#include <string.h>
#include <string>
#include <sstream>
#include <unordered_map>

#include "common/logging.h"
#include "gutil/dynamic_annotations.h"
#include "runtime/mem-pool.h"
#include "util/bit-util.h"

namespace impala {

/// Implementation of a free pool to recycle allocations. The pool is broken
/// up into 64 lists, one for each power of 2. Each allocation is rounded up
/// to the next power of 2 (or 8 bytes, whichever is larger). When the
/// allocation is freed, it is added to the corresponding free list.
/// Each allocation has an 8 byte header that immediately precedes the actual
/// allocation. If the allocation is owned by the user, the header contains
/// the ptr to the list that it should be added to on Free().
/// When the allocation is in the pool (i.e. available to be handed out), it
/// contains the link to the next allocation.
/// This has O(1) Allocate() and Free().
/// This is not thread safe.
/// TODO: consider integrating this with MemPool.
/// TODO: consider changing to something more granular than doubling.
class FreePool {
 public:
  /// C'tor, initializes the FreePool to be empty. All allocations come from the
  /// 'mem_pool'.
  FreePool(MemPool* mem_pool)
    : mem_pool_(mem_pool),
      net_allocations_(0) {
    memset(&lists_, 0, sizeof(lists_));
  }

  /// Allocates a buffer of size between [0, 2^62 - 1 - sizeof(FreeListNode)] bytes.
  uint8_t* Allocate(const int64_t requested_size) {
    DCHECK_GE(requested_size, 0);
    /// Return a non-nullptr dummy pointer. nullptr is reserved for failures.
    if (UNLIKELY(requested_size == 0)) return mem_pool_->EmptyAllocPtr();
    ++net_allocations_;
    /// MemPool allocations are 8-byte aligned, so making allocations < 8 bytes
    /// doesn't save memory and eliminates opportunities to recycle allocations.
    int64_t actual_size = std::max<int64_t>(8, requested_size);
    int free_list_idx = BitUtil::Log2Ceiling64(actual_size);
    DCHECK_LT(free_list_idx, NUM_LISTS);
    FreeListNode* allocation = lists_[free_list_idx].next;
    if (allocation == nullptr) {
      // There wasn't an existing allocation of the right size, allocate a new one.
      actual_size = 1LL << free_list_idx;
      allocation = reinterpret_cast<FreeListNode*>(
          mem_pool_->Allocate(actual_size + sizeof(FreeListNode)));
      if (UNLIKELY(allocation == nullptr)) {
        --net_allocations_;
        return nullptr;
      }
      // Memory will be returned unpoisoned from MemPool. Poison the whole range and then
      // deal with unpoisoning both allocation paths in one place below.
      ASAN_POISON_MEMORY_REGION(allocation, sizeof(FreeListNode) + actual_size);
    } else {
      // Remove this allocation from the list.
      lists_[free_list_idx].next = allocation->next;
    }
    DCHECK(allocation != nullptr);
#ifdef ADDRESS_SANITIZER
    uint8_t* ptr = reinterpret_cast<uint8_t*>(allocation);
    ASAN_UNPOISON_MEMORY_REGION(ptr, sizeof(FreeListNode) + requested_size);
    alloc_to_size_[ptr + sizeof(FreeListNode)] = requested_size;
#endif
    // Set the back node to point back to the list it came from so know where
    // to add it on Free().
    allocation->list = &lists_[free_list_idx];
    return reinterpret_cast<uint8_t*>(allocation) + sizeof(FreeListNode);
  }

  void Free(uint8_t* ptr) {
    if (UNLIKELY(ptr == nullptr || ptr == mem_pool_->EmptyAllocPtr())) return;
    --net_allocations_;
    FreeListNode* node = reinterpret_cast<FreeListNode*>(ptr - sizeof(FreeListNode));
    FreeListNode* list = node->list;
#ifndef NDEBUG
    CheckValidAllocation(list, ptr);
#endif
    // Add node to front of list.
    node->next = list->next;
    list->next = node;

#ifdef ADDRESS_SANITIZER
    int bucket_idx = (list - &lists_[0]);
    DCHECK_LT(bucket_idx, NUM_LISTS);
    int64_t allocation_size = 1LL << bucket_idx;
    ASAN_POISON_MEMORY_REGION(ptr, allocation_size);
    alloc_to_size_.erase(ptr);
#endif
  }

  /// Returns an allocation that is at least 'size'. If the current allocation backing
  /// 'ptr' is big enough, 'ptr' is returned. Otherwise a new one is made and the contents
  /// of ptr are copied into it.
  ///
  /// nullptr will be returned on allocation failure. It's the caller's responsibility to
  /// free the memory buffer pointed to by "ptr" in this case.
  uint8_t* Reallocate(uint8_t* ptr, int64_t size) {
    if (UNLIKELY(ptr == nullptr || ptr == mem_pool_->EmptyAllocPtr())) return Allocate(size);
    FreeListNode* node = reinterpret_cast<FreeListNode*>(ptr - sizeof(FreeListNode));
    FreeListNode* list = node->list;
#ifndef NDEBUG
    CheckValidAllocation(list, ptr);
#endif
    int bucket_idx = (list - &lists_[0]);
    DCHECK_LT(bucket_idx, NUM_LISTS);
    // This is the actual size of ptr.
    int64_t allocation_size = 1LL << bucket_idx;

    // If it's already big enough, just return the ptr.
    if (allocation_size >= size) {
      // Ensure that only first size bytes are unpoisoned. Need to poison whole region
      // first in case size is smaller than original allocation's size.
#ifdef ADDRESS_SANITIZER
      DCHECK(alloc_to_size_.find(ptr) != alloc_to_size_.end());
      int64_t prev_allocation_size = alloc_to_size_[ptr];
      if (prev_allocation_size > size) {
        // Allocation is shrinking: poison the 'freed' bytes.
        ASAN_POISON_MEMORY_REGION(ptr + size, prev_allocation_size - size);
      } else {
        // Allocation is growing: unpoison the newly allocated bytes.
        ASAN_UNPOISON_MEMORY_REGION(
            ptr + prev_allocation_size, size - prev_allocation_size);
      }
      alloc_to_size_[ptr] = size;
#endif
      return ptr;
    }

    // Make a new one. Since Allocate() already rounds up to powers of 2, this effectively
    // doubles for the caller.
    uint8_t* new_ptr = Allocate(size);
    if (LIKELY(new_ptr != nullptr)) {
#ifdef ADDRESS_SANITIZER
      DCHECK(alloc_to_size_.find(ptr) != alloc_to_size_.end());
      // Unpoison the region so that we can copy the old allocation to the new one.
      ASAN_UNPOISON_MEMORY_REGION(ptr, allocation_size);
#endif
      memcpy(new_ptr, ptr, allocation_size);
      Free(ptr);
    }
    return new_ptr;
  }

  MemTracker* mem_tracker() { return mem_pool_->mem_tracker(); }
  int64_t net_allocations() const { return net_allocations_; }

 private:
  static const int NUM_LISTS = 64;

  struct FreeListNode {
    /// Union for clarity when manipulating the node.
    union {
      FreeListNode* next; // Used when it is in the free list
      FreeListNode* list; // Used when it is being used by the caller.
    };
  };

  void CheckValidAllocation(FreeListNode* computed_list_ptr, uint8_t* allocation) const {
    // On debug, check that list is valid.
    bool found = false;
    for (int i = 0; i < NUM_LISTS && !found; ++i) {
      if (computed_list_ptr == &lists_[i]) found = true;
    }
    DCHECK(found) << "Could not find list for ptr: "
                  << reinterpret_cast<void*>(allocation)
                  << ". Allocation could have already been freed." << std::endl
                  << DebugString();
  }

  std::string DebugString() const {
    std::stringstream ss;
    ss << "FreePool: " << this << std::endl;
    for (int i = 0; i < NUM_LISTS; ++i) {
      FreeListNode* n = lists_[i].next;
      if (n == nullptr) continue;
      ss << i << ": ";
      while (n != nullptr) {
        uint8_t* ptr = reinterpret_cast<uint8_t*>(n);
        ptr += sizeof(FreeListNode);
        ss << reinterpret_cast<void*>(ptr);
        n = n->next;
        if (n != nullptr) ss << "->";
      }
      ss << std::endl;
    }
    return ss.str();
  }

  /// MemPool to allocate from. Unowned.
  MemPool* mem_pool_;

  /// One list head for each allocation size indexed by the LOG_2 of the allocation size.
  /// While it doesn't make too much sense to use this for very small (e.g. 8 byte)
  /// allocations, it makes the indexing easy.
  FreeListNode lists_[NUM_LISTS];

#ifdef ADDRESS_SANITIZER
  // For ASAN only: keep track of used bytes for each allocation (not including
  // FreeListNode header for each chunk). This allows us to poison each byte in a chunk
  // exactly once when reallocating from an existing chunk, rather than having to poison
  // the entire chunk each time.
  std::unordered_map<uint8_t*, int64_t> alloc_to_size_;
#endif

  /// Diagnostic counter that tracks (# Allocates - # Frees)
  int64_t net_allocations_;
};

}

#endif