// Copyright 2018 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 <lib/fit/defer.h>
#include <lib/fzl/pinned-vmo.h>

#include <climits>
#include <limits>
#include <memory>

namespace fzl {

zx_status_t PinnedVmo::Pin(const zx::vmo& vmo, const zx::bti& bti, uint32_t options) {
  if (!vmo.is_valid()) {
    return ZX_ERR_INVALID_ARGS;
  }
  // To pin the entire VMO, we need to get the length:
  zx_status_t res;
  uint64_t vmo_size;
  res = vmo.get_size(&vmo_size);
  if (res != ZX_OK) {
    return res;
  }
  return PinInternal(0, vmo_size, vmo, bti, options);
}

zx_status_t PinnedVmo::PinRange(uint64_t offset, uint64_t len, const zx::vmo& vmo,
                                const zx::bti& bti, uint32_t options) {
  const size_t kPageSize = zx_system_get_page_size();
  if ((len & (kPageSize - 1)) || (offset & (kPageSize - 1)) || len == 0) {
    return ZX_ERR_INVALID_ARGS;
  }

  return PinInternal(offset, len, vmo, bti, options);
}

zx_status_t PinnedVmo::PinInternal(uint64_t offset, uint64_t len, const zx::vmo& vmo,
                                   const zx::bti& bti, uint32_t options) {
  const size_t kPageSize = zx_system_get_page_size();
  zx_status_t res;

  // If we are holding a pinned memory token, then we are already holding a
  // pinned VMO.  It is an error to try and pin a new VMO without first
  // explicitly unpinning the old one.
  if (pmt_.is_valid()) {
    ZX_DEBUG_ASSERT(regions_ != nullptr);
    ZX_DEBUG_ASSERT(region_count_ > 0);
    return ZX_ERR_BAD_STATE;
  }

  // Check our args, read/write/bti_contiguous is all that users may ask for.
  constexpr uint32_t kAllowedOptions = ZX_BTI_PERM_READ | ZX_BTI_PERM_WRITE | ZX_BTI_CONTIGUOUS;
  if (((options & kAllowedOptions) != options) || !vmo.is_valid() || !bti.is_valid()) {
    return ZX_ERR_INVALID_ARGS;
  }

  // Allocate storage for the results.
  ZX_DEBUG_ASSERT((len > 0) && !(len & (kPageSize - 1)));
  ZX_DEBUG_ASSERT((len / kPageSize) < std::numeric_limits<uint32_t>::max());
  ZX_DEBUG_ASSERT(!(offset & (kPageSize - 1)));
  fbl::AllocChecker ac;
  uint32_t page_count = static_cast<uint32_t>(len / kPageSize);
  if (options & ZX_BTI_CONTIGUOUS) {
    page_count = 1;
  }

  std::unique_ptr<zx_paddr_t[]> addrs(new (&ac) zx_paddr_t[page_count]);
  if (!ac.check()) {
    return ZX_ERR_NO_MEMORY;
  }

  // Now actually pin the region.
  res = bti.pin(options, vmo, offset, len, addrs.get(), page_count, &pmt_);
  if (res != ZX_OK) {
    return res;
  }

  // From here on out, if anything goes wrong, we need to make sure to clean
  // up.  Setup an autocall to take care of this for us.
  auto cleanup = fit::defer([&]() { UnpinInternal(); });

  if (options & ZX_BTI_CONTIGUOUS) {
    // We can do less work in the contiguous case.
    regions_.reset(new (&ac) Region[1]);
    if (!ac.check()) {
      return ZX_ERR_NO_MEMORY;
    }

    region_count_ = 1;
    regions_[0].phys_addr = addrs[0];
    // The region is the size of the entire vmo in the contiguous case.
    regions_[0].size = len;
    cleanup.cancel();
    return ZX_OK;
  }

  // Do a quick pass over the pages to figure out how many adjacent pages we
  // can merge.  This will let us know how many regions we will need storage
  // for our regions array.
  zx_paddr_t last = addrs[0];
  region_count_ = 1;
  for (uint32_t i = 1; i < page_count; ++i) {
    if (addrs[i] != (last + kPageSize)) {
      ++region_count_;
    }
    last = addrs[i];
  }

  // Allocate storage for our regions.
  regions_.reset(new (&ac) Region[region_count_]);
  if (!ac.check()) {
    return ZX_ERR_NO_MEMORY;
  }

  // Finally, go ahead and merge any adjacent pages to compute our set of
  // regions and we should be good to go;
  regions_[0].phys_addr = addrs[0];
  regions_[0].size = kPageSize;
  for (uint32_t i = 1, j = 0; i < page_count; ++i) {
    ZX_DEBUG_ASSERT(j < region_count_);

    if ((regions_[j].phys_addr + regions_[j].size) == addrs[i]) {
      // Merge!
      regions_[j].size += kPageSize;
    } else {
      // New Region!
      ++j;
      ZX_DEBUG_ASSERT(j < region_count_);
      regions_[j].phys_addr = addrs[i];
      regions_[j].size = kPageSize;
    }
  }

  cleanup.cancel();
  return ZX_OK;
}

void PinnedVmo::Unpin() {
  if (!pmt_.is_valid()) {
    ZX_DEBUG_ASSERT(regions_ == nullptr);
    ZX_DEBUG_ASSERT(region_count_ == 0);
    return;
  }

  ZX_DEBUG_ASSERT(regions_ != nullptr);
  ZX_DEBUG_ASSERT(region_count_ > 0);

  UnpinInternal();
}

void PinnedVmo::UnpinInternal() {
  ZX_DEBUG_ASSERT(pmt_.is_valid());

  // Given the level of sanity checking we have done so far, it should be
  // completely impossible for us to fail to unpin this memory.
  [[maybe_unused]] zx_status_t res;
  res = pmt_.unpin();
  ZX_DEBUG_ASSERT(res == ZX_OK);

  regions_.reset();
  region_count_ = 0;
}

}  // namespace fzl