// Copyright 2016 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <string.h>

#include <fbl/algorithm.h>
#include <fbl/alloc_checker.h>
#include <region-alloc/region-alloc.h>

// Support for Pool allocated bookkeeping
RegionAllocator::RegionPool::RefPtr RegionAllocator::RegionPool::Create(size_t max_memory) {
  // Sanity check our allocation arguments.
  if (SLAB_SIZE > max_memory) {
    return nullptr;
  }

  fbl::AllocChecker ac;
  auto ret = fbl::AdoptRef(new (&ac) RegionPool(max_memory / SLAB_SIZE));

  if (!ac.check()) {
    return nullptr;
  }

  return ret;
}

RegionAllocator::~RegionAllocator() {
  // No one should be destroying us while we have allocations in flight.
  ZX_DEBUG_ASSERT(allocated_regions_by_base_.is_empty());

  // We should have the same number of regions sorted by base address and by
  // size.
  ZX_DEBUG_ASSERT(avail_regions_by_base_.size() == avail_regions_by_size_.size());

  // We have to have a region pool assigned to us, or our available regions
  // need to be empty.
  ZX_DEBUG_ASSERT((region_pool_ != nullptr) ||
                  (avail_regions_by_base_.is_empty() && avail_regions_by_size_.is_empty()));

  // Return all of our bookkeeping to our region pool.
  avail_regions_by_base_.clear();
  while (!avail_regions_by_size_.is_empty()) {
    region_pool_->Delete(avail_regions_by_size_.pop_front());
  }
}

void RegionAllocator::Reset() {
  fbl::AutoLock alloc_lock(&alloc_lock_);

  ZX_DEBUG_ASSERT((region_pool_ != nullptr) || avail_regions_by_base_.is_empty());

  Region* removed;
  while ((removed = avail_regions_by_base_.pop_front()) != nullptr) {
    avail_regions_by_size_.erase(*removed);
    region_pool_->Delete(removed);
  }

  ZX_DEBUG_ASSERT(avail_regions_by_base_.is_empty());
  ZX_DEBUG_ASSERT(avail_regions_by_size_.is_empty());
}

zx_status_t RegionAllocator::SetRegionPool(const RegionPool::RefPtr& region_pool) {
  fbl::AutoLock alloc_lock(&alloc_lock_);

  if (!allocated_regions_by_base_.is_empty() || !avail_regions_by_base_.is_empty())
    return ZX_ERR_BAD_STATE;

  region_pool_ = region_pool;
  return ZX_OK;
}

zx_status_t RegionAllocator::AddRegion(const ralloc_region_t& region, bool allow_overlap) {
  fbl::AutoLock alloc_lock(&alloc_lock_);

  // Start with sanity checks
  zx_status_t ret = AddSubtractSanityCheckLocked(region);
  if (ret != ZX_OK)
    return ret;

  // Make sure that we do not intersect with the available regions if we do
  // not allow overlaps.
  if (!allow_overlap && IntersectsLocked(avail_regions_by_base_, region))
    return ZX_ERR_INVALID_ARGS;

  // All sanity checks passed.  Grab a piece of free bookkeeping from our pool,
  // fill it out, then add it to the sets of available regions (indexed by
  // base address as well as size)
  Region* to_add = region_pool_->New(this);
  if (to_add == nullptr)
    return ZX_ERR_NO_MEMORY;

  to_add->base = region.base;
  to_add->size = region.size;

  AddRegionToAvailLocked(to_add, allow_overlap);

  return ZX_OK;
}

zx_status_t RegionAllocator::SubtractRegion(const ralloc_region_t& to_subtract,
                                            bool allow_incomplete) {
  fbl::AutoLock alloc_lock(&alloc_lock_);

  // Start with sanity checks
  zx_status_t ret = AddSubtractSanityCheckLocked(to_subtract);
  if (ret != ZX_OK)
    return ret;

  // Make a copy of the region to subtract.  We may need to modify the region
  // as part of the subtraction algorithm.
  ralloc_region_t region = to_subtract;

  // Find the region whose base address is <= the specified region (if any).
  // If we do not allow incomplete subtraction, this is the region which must
  // entirely contain the subtracted region.
  //
  // Additionally, the only time we should ever need any extra bookkeeping
  // allocation is if this region (if present) needs to be split into two
  // regions.
  auto before = avail_regions_by_base_.upper_bound(region.base);
  auto after = before--;
  uint64_t region_end = region.base + region.size;  // exclusive end
  uint64_t before_end = 0;

  if (before.IsValid()) {
    before_end = before->base + before->size;  // exclusive end
    if ((region.base >= before->base) && (region_end <= before_end)) {
      // Looks like we found an available region which completely contains the
      // region to be subtracted.  Handle the 4 possible cases.

      // Case 1: The regions are the same.  This one is easy.
      if ((region.base == before->base) && (region_end == before_end)) {
        Region* removed = avail_regions_by_base_.erase(before);
        avail_regions_by_size_.erase(*removed);

        ZX_DEBUG_ASSERT(region_pool_ != nullptr);
        region_pool_->Delete(removed);

        return ZX_OK;
      }

      // Case 2: before completely contains region.  The before region needs
      // to be split into two regions.  If we are out of bookkeeping space, we
      // are out of luck.
      if ((region.base != before->base) && (region_end != before_end)) {
        Region* second = region_pool_->New(this);
        if (second == nullptr)
          return ZX_ERR_NO_MEMORY;

        // Looks like we have the memory we need.  Compute the base/size of
        // the two regions which will be left over, then update the first
        // region's position in the size index, and add the second region to
        // the set of available regions.
        Region* first = avail_regions_by_size_.erase(*before);
        first->size = region.base - first->base;
        second->base = region_end;
        second->size = before_end - region_end;

        avail_regions_by_size_.insert(first);
        avail_regions_by_base_.insert(second);
        avail_regions_by_size_.insert(second);
        return ZX_OK;
      }

      // Case 3: region trims the front of before.  Update before's base and
      // size, and recompute its position in the size index.  Note: there is
      // no need to recompute its position in the base index, this has not
      // changed.
      if (region.base == before->base) {
        ZX_DEBUG_ASSERT(region_end < before_end);

        Region* bptr = avail_regions_by_size_.erase(*before);
        bptr->base += region.size;
        bptr->size -= region.size;
        avail_regions_by_size_.insert(bptr);

        return ZX_OK;
      }

      // Case 4: region trims the end of before.  Update before's size and
      // recompute its position in the size index.
      ZX_DEBUG_ASSERT(region.base != before->base);
      ZX_DEBUG_ASSERT(region_end == before_end);

      Region* bptr = avail_regions_by_size_.erase(*before);
      bptr->size -= region.size;
      avail_regions_by_size_.insert(bptr);

      return ZX_OK;
    }
  }

  // If we have gotten this far, then there is no single region in the
  // available set which completely contains the subtraction region.  We
  // cannot continue unless allow_incomplete is true.
  if (!allow_incomplete)
    return ZX_ERR_INVALID_ARGS;

  // Great!  At this point we know that we are going to succeed, we just need
  // to go about updating all of the bookkeeping.  We may need to trim the end
  // of the region which comes before us, and then consume some of the regions
  // which come after us, finishing by trimming the front of (at most) one of
  // the regions which comes after us.  At no point in time do we need to
  // allocate any more bookkeeping, success is guaranteed.  Start by
  // considering the before region.
  if (before.IsValid()) {
    ZX_DEBUG_ASSERT(region.base >= before->base);
    ZX_DEBUG_ASSERT(region_end > before_end);
    if (before_end > region.base) {
      // No matter what, 'before' needs to be removed from the size index.
      Region* bptr = avail_regions_by_size_.erase(*before);

      // If before's base is the same as the region's base, then we are
      // subtracting out all of before.  Otherwise, we are trimming the back
      // before and need to recompute its size and position in the size index.
      if (bptr->base == region.base) {
        avail_regions_by_base_.erase(*bptr);

        ZX_DEBUG_ASSERT(region_pool_ != nullptr);
        region_pool_->Delete(bptr);
      } else {
        bptr->size = region.base - bptr->base;
        avail_regions_by_size_.insert(bptr);
      }

      // Either way, the region we are subtracting now starts where before
      // used to end.
      region.base = before_end;
      region.size = region_end - region.base;
      ZX_DEBUG_ASSERT(region.size > 0);
    }
  }

  // While there are regions whose base address comes after the base address
  // of what we want to subtract, we need to do one of three things...
  //
  // 1) Consume entire regions which are contained entirely within our
  //    subtraction region.
  // 2) Trim the front of a region which is clipped by our subtraction region.
  // 3) Stop because all remaining regions start after the end of our
  //    subtraction region.
  while (after.IsValid()) {
    ZX_DEBUG_ASSERT(after->base > region.base);

    // Case #3
    if (after->base >= region_end)
      break;

    // Cases #1 and #2.  No matter what, we need to...
    // 1) Advance after, re-naming the old 'after' to 'trim in the process.
    // 2) Remove trim from the size index.
    auto trim_iter = after++;
    Region* trim = avail_regions_by_size_.erase(*trim_iter);
    uint64_t trim_end = trim->base + trim->size;

    if (trim_end > region_end) {
      // Case #2.  We are guaranteed to be done at this point.
      trim->base = region_end;
      trim->size = trim_end - trim->base;
      avail_regions_by_size_.insert(trim);
      break;
    }

    // Case #1.  Advance the subtraction region to the end of the region we
    // just trimmed into oblivion.  If we have run out of region to trim,
    // then we know we are done.
    avail_regions_by_base_.erase(*trim);
    region.base = trim_end;
    region.size = region_end - region.base;

    ZX_DEBUG_ASSERT(region_pool_ != nullptr);
    region_pool_->Delete(trim);

    if (!region.size)
      break;
  }

  // Sanity check.  The number of elements in the base index should match the
  // number of elements in the size index.
  ZX_DEBUG_ASSERT(avail_regions_by_base_.size() == avail_regions_by_size_.size());
  return ZX_OK;
}

zx_status_t RegionAllocator::GetRegion(uint64_t size, uint64_t alignment,
                                       Region::UPtr& out_region) {
  fbl::AutoLock alloc_lock(&alloc_lock_);

  // Check our RegionPool
  if (region_pool_ == nullptr)
    return ZX_ERR_BAD_STATE;

  // Sanity check the arguments.
  out_region = nullptr;
  if (!size || !alignment || !fbl::is_pow2(alignment))
    return ZX_ERR_INVALID_ARGS;

  // Compute the things we will need round-up align base addresses.
  uint64_t mask = alignment - 1;
  uint64_t inv_mask = ~mask;

  // Start by using our size index to look up the first available region which
  // is large enough to hold this allocation (if any)
  auto iter = avail_regions_by_size_.lower_bound({.base = 0, .size = size});

  // Consider all of the regions which are large enough to hold our
  // allocation.  Stop as soon as we find one which can satisfy the alignment
  // restrictions.
  uint64_t aligned_base = 0;
  while (iter.IsValid()) {
    ZX_DEBUG_ASSERT(iter->size >= size);
    aligned_base = (iter->base + mask) & inv_mask;
    uint64_t overhead = aligned_base - iter->base;
    uint64_t leftover = iter->size - size;

    // We have a usable region if the aligned base address has not wrapped
    // the address space, and if overhead required to align the allocation
    // is not larger than what is leftover in the region after performing
    // the allocation.
    if ((aligned_base >= iter->base) && (overhead <= leftover))
      break;

    ++iter;
  }

  if (!iter.IsValid())
    return ZX_ERR_NOT_FOUND;

  return AllocFromAvailLocked(iter, out_region, aligned_base, size);
}

zx_status_t RegionAllocator::GetRegion(const ralloc_region_t& requested_region,
                                       Region::UPtr& out_region) {
  fbl::AutoLock alloc_lock(&alloc_lock_);

  // Check our RegionPool
  if (region_pool_ == nullptr)
    return ZX_ERR_BAD_STATE;

  uint64_t base = requested_region.base;
  uint64_t size = requested_region.size;

  // Sanity check the arguments.
  out_region = nullptr;
  if (!size || ((base + size) < base))
    return ZX_ERR_INVALID_ARGS;

  // Find the first available region whose base address is strictly greater
  // than the one we are looking for, then back up one.
  auto iter = avail_regions_by_base_.upper_bound(base);
  --iter;

  // If the iterator is invalid, then we cannot satisfy this request.  If it
  // is valid, then we can satisfy this request if and only if the region we
  // found completely contains the requested region.
  if (!iter.IsValid())
    return ZX_ERR_NOT_FOUND;

  // We know that base must be >= iter->base
  // We know that iter->size is non-zero.
  // Therefore, we know that base is in the range [iter.start, iter.end]
  // We know that request.end is > base
  // Therefore request.end is > iter.base
  //
  // So, if request.end <= iter.end, we know that request is completely
  // contained within iter.  It does not matter if we use the inclusive or
  // exclusive end to check, as long as we are consistent.
  ZX_DEBUG_ASSERT(iter->size > 0);
  ZX_DEBUG_ASSERT(iter->base <= base);
  uint64_t req_end = base + size - 1;
  uint64_t iter_end = iter->base + iter->size - 1;
  if (req_end > iter_end)
    return ZX_ERR_NOT_FOUND;

  // Great, we have found a region which should be able to satisfy our
  // allocation request.  Get an iterator for the by-size index, then use the
  // common AllocFromAvailLocked method to handle the bookkeeping involved.
  auto by_size_iter = avail_regions_by_size_.make_iterator(*iter);
  return AllocFromAvailLocked(by_size_iter, out_region, base, size);
}

zx_status_t RegionAllocator::AddSubtractSanityCheckLocked(const ralloc_region_t& region) {
  // Check our RegionPool
  if (region_pool_ == nullptr)
    return ZX_ERR_BAD_STATE;

  // Sanity check the region to make sure that it is well formed.  We do not
  // allow a region which is of size zero, or which wraps around the
  // allocation space.
  if ((region.base + region.size) <= region.base)
    return ZX_ERR_INVALID_ARGS;

  // Next, make sure the region we are adding or subtracting does not
  // intersect any region which is currently allocated.
  if (IntersectsLocked(allocated_regions_by_base_, region))
    return ZX_ERR_INVALID_ARGS;

  return ZX_OK;
}

void RegionAllocator::ReleaseRegion(Region* region) {
  fbl::AutoLock alloc_lock(&alloc_lock_);

  ZX_DEBUG_ASSERT(region != nullptr);

  // When a region comes back from a user, it should be in the
  // allocated_regions_by_base tree, but not in either of the avail_regions
  // trees, and not in any free list.  Remove it from the allocated_regions
  // bookkeeping and add it back to the available regions.
  ZX_DEBUG_ASSERT(fbl::InContainer<SortByBaseTag>(*region));
  ZX_DEBUG_ASSERT(!fbl::InContainer<SortBySizeTag>(*region));

  allocated_regions_by_base_.erase(*region);
  AddRegionToAvailLocked(region);
}

zx_status_t RegionAllocator::AllocFromAvailLocked(Region::WAVLTreeSortBySize::iterator source,
                                                  Region::UPtr& out_region, uint64_t base,
                                                  uint64_t size) {
  ZX_DEBUG_ASSERT(out_region == nullptr);
  ZX_DEBUG_ASSERT(source.IsValid());
  ZX_DEBUG_ASSERT(base >= source->base);
  ZX_DEBUG_ASSERT(size <= source->size);

  uint64_t overhead = base - source->base;
  ZX_DEBUG_ASSERT(overhead < source->size);

  uint64_t leftover = source->size - size;
  ZX_DEBUG_ASSERT(leftover >= overhead);

  // Great, we found a region.  We may have to split the available region into
  // up to 2 sub regions depending on where the aligned allocation lies in the
  // region.  Figure out how much splitting we need to do and attempt to
  // allocate the bookkeeping.
  bool split_before = base != source->base;
  bool split_after = overhead < leftover;

  if (!split_before && !split_after) {
    // If no splits are required, then this should be easy.  Take the region
    // out of the avail bookkeeping, add it to the allocated bookkeeping and
    // we are finished.
    Region* region = avail_regions_by_size_.erase(source);
    avail_regions_by_base_.erase(*region);
    allocated_regions_by_base_.insert(region);
    out_region.reset(region);
  } else if (!split_before) {
    // If we only have to split after, then this region is aligned with what
    // we want to allocate, but we will not use all of it.  Break it into
    // two pieces and return the one which comes first.
    Region* before_region = region_pool_->New(this);
    if (before_region == nullptr)
      return ZX_ERR_NO_MEMORY;

    Region* after_region = avail_regions_by_size_.erase(source);

    before_region->base = after_region->base;
    before_region->size = size;
    after_region->base += size;
    after_region->size -= size;

    avail_regions_by_size_.insert(after_region);
    allocated_regions_by_base_.insert(before_region);

    out_region.reset(before_region);
  } else if (!split_after) {
    // If we only have to split before, then this region is not aligned
    // properly with what we want to allocate, but we will use the entire
    // region (after aligning).  Break it into two pieces and return the one
    // which comes after.
    Region* after_region = region_pool_->New(this);
    if (after_region == nullptr)
      return ZX_ERR_NO_MEMORY;

    Region* before_region = avail_regions_by_size_.erase(source);

    after_region->base = base;
    after_region->size = size;
    before_region->size -= size;

    avail_regions_by_size_.insert(before_region);
    allocated_regions_by_base_.insert(after_region);

    out_region.reset(after_region);
  } else {
    // Looks like we need to break our region into 3 chunk and return the
    // middle chunk.  Start by grabbing the bookkeeping we require first.
    Region* region = region_pool_->New(this);
    if (region == nullptr)
      return ZX_ERR_NO_MEMORY;

    Region* after_region = region_pool_->New(this);
    if (after_region == nullptr) {
      ZX_DEBUG_ASSERT(region_pool_ != nullptr);
      region_pool_->Delete(region);
      return ZX_ERR_NO_MEMORY;
    }

    Region* before_region = avail_regions_by_size_.erase(source);

    region->base = before_region->base + overhead;
    region->size = size;
    after_region->base = region->base + region->size;
    after_region->size = before_region->size - size - overhead;
    before_region->size = overhead;

    avail_regions_by_size_.insert(before_region);
    avail_regions_by_size_.insert(after_region);
    avail_regions_by_base_.insert(after_region);
    allocated_regions_by_base_.insert(region);

    out_region.reset(region);
  }
  return ZX_OK;
}

void RegionAllocator::AddRegionToAvailLocked(Region* region, bool allow_overlap) {
  // Sanity checks.  This region should not exist in any bookkeeping, and
  // should not overlap with any of the regions we are currently tracking.
  ZX_DEBUG_ASSERT(!fbl::InContainer<SortByBaseTag>(*region));
  ZX_DEBUG_ASSERT(!fbl::InContainer<SortBySizeTag>(*region));
  ZX_DEBUG_ASSERT(!IntersectsLocked(allocated_regions_by_base_, *region));
  ZX_DEBUG_ASSERT(allow_overlap || !IntersectsLocked(avail_regions_by_base_, *region));

  // Find the region which comes before us and the region which comes after us
  // in the tree.
  auto before = avail_regions_by_base_.upper_bound(region->base);
  auto after = before--;

  // Merge with the region which comes before us if we can.
  uint64_t region_end = (region->base + region->size);  // exclusive end
  if (before.IsValid()) {
    ZX_DEBUG_ASSERT(before->base <= region->base);

    uint64_t before_end = (before->base + before->size);  // exclusive end
    if (allow_overlap ? (before_end >= region->base) : (before_end == region->base)) {
      region_end = fbl::max(region_end, before_end);
      region->base = before->base;

      auto removed = avail_regions_by_base_.erase(before);
      avail_regions_by_size_.erase(*removed);
      ZX_DEBUG_ASSERT(region_pool_ != nullptr);
      region_pool_->Delete(removed);
    }
  }

  // Merge with the region which comes after us if we can, keep merging if we
  // allow overlaps.
  while (after.IsValid()) {
    ZX_DEBUG_ASSERT(region->base < after->base);

    if (!(allow_overlap ? (region_end >= after->base) : (region_end == after->base)))
      break;

    uint64_t after_end = (after->base + after->size);
    region_end = fbl::max(region_end, after_end);

    auto remove_me = after++;
    auto removed = avail_regions_by_base_.erase(remove_me);
    avail_regions_by_size_.erase(*removed);
    ZX_DEBUG_ASSERT(region_pool_ != nullptr);
    region_pool_->Delete(removed);

    if (!allow_overlap)
      break;
  }

  // Update the region's size to reflect any mergers which may have taken
  // place, then add the region to the two indexes.
  region->size = region_end - region->base;
  avail_regions_by_base_.insert(region);
  avail_regions_by_size_.insert(region);
}

bool RegionAllocator::IntersectsLocked(const Region::WAVLTreeSortByBase& tree,
                                       const ralloc_region_t& region) {
  // Find the first entry in the tree whose base is >= region.base.  If this
  // element exists, and its base is < the exclusive end of region, then
  // we have an intersection.
  auto iter = tree.lower_bound(region.base);
  if (iter.IsValid() && (iter->base < (region.base + region.size)))
    return true;

  // Check the element before us in the tree.  If it exists, we know that it's
  // base is < region.base.  If it's exclusive end is >= region.base, then we
  // have an intersection.
  --iter;
  if (iter.IsValid() && (region.base < (iter->base + iter->size)))
    return true;

  return false;
}