// Copyright 2019 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 <stddef.h>

#include <fbl/algorithm.h>
#include <fbl/alloc_checker.h>
#include <id_allocator/id_allocator.h>
#include <lib/zircon-internal/debug.h>
#include <zircon/errors.h>
#include <zircon/types.h>

#define ZXDEBUG 0

namespace id_allocator {

namespace {
constexpr size_t kMaxId = SIZE_MAX;
constexpr int8_t kLevelBits = 6;
constexpr size_t kLevelMask = static_cast<size_t>((1 << kLevelBits) - 1);
constexpr size_t kMaxChildren = static_cast<size_t>(1 << kLevelBits);

constexpr int8_t log64_int8_ceil(size_t id_count) {
    size_t max = 1 << kLevelBits;
    int8_t i = 1;

    while (max < id_count) {
        max = max << kLevelBits;
        i++;
    }

    return i;
}
// Returns number of levels needed to map id_count ids
// id_count <= 64, levels == 1
// id_count <= 64 * 64, levels = 2
constexpr int8_t GetLevels(size_t id_count) {
    return log64_int8_ceil(id_count);
}

} // namespace

size_t IdAllocator::LevelBitCount(int8_t level) const {
    if (level) {
        size_t pad = 1;
        pad = (pad << (kLevelBits * level)) - 1;
        return ((id_count_ + pad) >> (kLevelBits * level));
    }

    return id_count_;
}

size_t IdAllocator::LevelBitCountRounded(int8_t level) const {
    return fbl::round_up(LevelBitCount(level), kMaxChildren);
}

bool IdAllocator::SetBitAt(int8_t level, size_t bit) {
    ZX_ASSERT(bit < levels_[level].size());
    levels_[level].SetOne(bit);
    bit = bit & ~kLevelMask;
    return levels_[level].Get(bit, bit + kMaxChildren);
}

void IdAllocator::ClearBitAt(int8_t level, size_t bit) {
    ZX_ASSERT(bit < levels_[level].size());
    levels_[level].ClearOne(bit);
}

size_t IdAllocator::FindFirstUnset(int8_t level, size_t base_index) const {
    size_t first_zero;

    ZX_ASSERT((base_index & kLevelMask) == 0);
    if (!levels_[level].Get(base_index, base_index + kMaxChildren, &first_zero)) {
        return first_zero;
    }
    return kMaxId;
}

size_t IdAllocator::Find() const {
    size_t id = 0, index;
    for (int8_t level = static_cast<int8_t>(level_count_ - 1); level >= 0; level--) {
        id = id << kLevelBits;
        index = FindFirstUnset(level, id);
        if (index == kMaxId) {
            // if all ids are busy then all bits at root level will be set.
            ZX_ASSERT(level == level_count_ - 1);
            return kMaxId;
        }
        id = id | index;
    }

    return id;
}

void IdAllocator::MarkBusyInternal(size_t id, int8_t level) {
    ZX_ASSERT(level < level_count_);
    size_t index = id;
    bool all_children_busy;
    for (; level < level_count_; level++) {
        all_children_busy = SetBitAt(level, index);
        if (all_children_busy == false) {
            break;
        }
        index = index >> kLevelBits;
    }
}

void IdAllocator::MarkBusy(size_t id) {
    MarkBusyInternal(id, 0);
}

void IdAllocator::MarkFreeInternal(size_t id, int8_t level) {
    size_t index = id;
    for (; level < level_count_; level++) {
        if (level != 0) {
            if (levels_[level].GetOne(index) == false) {
                break;
            }
        }

        ClearBitAt(level, index);
        index = index >> kLevelBits;
    }
}

void IdAllocator::MarkFree(size_t id) {
    MarkFreeInternal(id, 0);
}

zx_status_t IdAllocator::Allocate(size_t* out) {
    size_t id;

    if ((id = Find()) >= id_count_) {
        xprintf("Found an id:%ld level:%d id_count:%ld\n", id, level_count_, id_count_);
        return ZX_ERR_NO_RESOURCES;
    }
    xprintf("Setting id:%ld level:%d id_count:%ld\n", id, level_count_, id_count_);
    ZX_ASSERT(id < id_count_);

    ZX_ASSERT(!IsBusy(id));
    MarkBusy(id);

    *out = id;
    return ZX_OK;
}

bool IdAllocator::IsBusy(size_t id) const {
    if (id >= id_count_) {
        return false;
    }
    return levels_[0].GetOne(id);
}

zx_status_t IdAllocator::MarkAllocated(size_t id) {
    if (id >= id_count_) {
        return ZX_ERR_OUT_OF_RANGE;
    }
    if (IsBusy(id)) {
        return ZX_ERR_BAD_STATE;
    }
    xprintf("Setting id:%ld level:%d id_count:%ld\n", id, level_count_, id_count_);
    MarkBusy(id);
    return ZX_OK;
}

zx_status_t IdAllocator::Free(size_t id) {
    if (id >= id_count_) {
        return ZX_ERR_OUT_OF_RANGE;
    }

    if (!IsBusy(id)) {
        return ZX_ERR_BAD_STATE;
    }
    xprintf("Freeing id:%ld level:%d id_count:%ld\n", id, level_count_, id_count_);
    MarkFree(id);
    return ZX_OK;
}

// Unallocatable bits are those sets of bits that are out of range for current
// value of id_count_. These bits are set so that searching for free bits becomes
// faster.
// For simplicity we iterate over each unallocatable bit. This isn't a lot given
// there can be at most 63 unallocatable bits for a level
// TODO(auradkar): Use range set
void IdAllocator::MarkUnallocatable(int8_t level) {
    size_t start = LevelBitCount(level);
    size_t end = LevelBitCountRounded(level);

    for (size_t i = start; i < end; i++) {
        MarkBusyInternal(i, level);
    }
}

// For simplicity we iterate over each unallocatable bit. This isn't a lot given
// there can be at most 63 unallocatable bits for a level
// TODO(auradkar): range clear
void IdAllocator::MarkAllAllocatable(int8_t level) {
    size_t start = LevelBitCount(level);
    size_t end = LevelBitCountRounded(level);

    for (size_t i = start; i < end; i++) {
        MarkFreeInternal(i, level);
    }
}

// Grow should not destroy existing data on failure. We leave successfully
// grown levels levels as is on failure. They are marked busy and
// unallocatable anyway
zx_status_t IdAllocator::GrowInternal(size_t id_count) {
    zx_status_t status;
    xprintf("Growing from %ld to %ld\n", id_count_, id_count);
    id_count_ = id_count;
    level_count_ = GetLevels(id_count_);
    for (int8_t level = 0; level < level_count_; level++) {
        size_t lsize = LevelBitCountRounded(level);
        // We are adding a new level
        if (levels_[level].size() == 0) {
            xprintf("Resetting level %d to size %ld\n", level, lsize);
            status = levels_[level].Reset(lsize);
            // If we are adding a new parent level and if all the children
            // bits are busy, we set the parent bit as busy.
            // When we add a new level, only one parent (at index 0) will
            // inherit all the existing children.
            if ((level > 0) && (levels_[level - 1].Get(0, kMaxChildren))) {
                levels_[level].SetOne(0);
            }
        } else {
            // We are extending an existing level
            xprintf("Growing level %d from %ld to %ld\n", level,
                    levels_[level].size(), lsize);
            status = levels_[level].Grow(lsize);
        }

        if (status != ZX_OK) {
            return status;
        }
        MarkUnallocatable(level);
    }
    return ZX_OK;
}

zx_status_t IdAllocator::Grow(size_t id_count) {
    size_t old_id_count = id_count_;
    if (id_count == id_count_) {
        return ZX_OK;
    }

    if (id_count < id_count_) {
        return ZX_ERR_INVALID_ARGS;
    }

    if (id_count >= kMaxId) {
        return ZX_ERR_OUT_OF_RANGE;
    }

    for (int8_t level = 0; level < level_count_; level++) {
        MarkAllAllocatable(level);
    }
    zx_status_t status = GrowInternal(id_count);
    if (status != ZX_OK) {
        // This should never fail. We should not have freed any
        // existing resources
        ZX_ASSERT(GrowInternal(old_id_count) == ZX_OK);
    }
    return status;
}

zx_status_t IdAllocator::Shrink(size_t id_count) {
    if (id_count == id_count_) {
        return ZX_OK;
    }

    if ((id_count > id_count_)) {
        return ZX_ERR_OUT_OF_RANGE;
    }
    int8_t old_level_count = GetLevels(id_count_);
    id_count_ = id_count;
    level_count_ = GetLevels(id_count);

    for (int8_t level = 0; level < old_level_count; level++) {
        // Free any level that was allocated but is not needed anymore
        if (level >= level_count_) {
            ZX_ASSERT(levels_[level].Reset(0) == ZX_OK);
            continue;
        }
        size_t lsize = LevelBitCountRounded(level);
        ZX_ASSERT(levels_[level].Shrink(lsize) == ZX_OK);
        MarkUnallocatable(level);
    }

    return ZX_OK;
}

zx_status_t IdAllocator::Reset(size_t id_count) {
    if ((id_count >= kMaxId)) {
        return ZX_ERR_OUT_OF_RANGE;
    }

    zx_status_t status;
    if (id_count < id_count_) {
        status = Shrink(id_count);
    } else {
        status = Grow(id_count);
    }

    if (status != ZX_OK) {
        return status;
    }

    for (int8_t level = 0; level < level_count_; level++) {
        ZX_ASSERT(levels_[level].Reset(levels_[level].size()) == ZX_OK);
        MarkUnallocatable(level);
    }

    return ZX_OK;
}

void IdAllocator::Dump() const {
    xprintf("kMaxLevels:%d id_count:%lu level_count:%d\n",
            kMaxLevels, id_count_, level_count_);
    for (int8_t level = static_cast<int8_t>(level_count_ - 1); level >= 0; level--) {
        xprintf("\nlevel: %d\n", level);
        for (size_t index = 0; index < levels_[level].size(); index++) {
            xprintf("%d", levels_[level].GetOne(index) ? 1 : 0);
        }
    }
    xprintf("\n");
}

zx_status_t IdAllocator::Create(size_t id_count,
                                std::unique_ptr<IdAllocator>* ida_out) {
    if (id_count >= kMaxId) {
        return ZX_ERR_OUT_OF_RANGE;
    }

    fbl::AllocChecker ac;
    IdAllocator* ida = new (&ac) IdAllocator;

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

    std::unique_ptr<IdAllocator> idap(ida);
    zx_status_t status;

    for (int8_t i = 0; i < kMaxLevels; i++) {
        if ((status = idap->levels_[i].Reset(0)) != ZX_OK) {
            return status;
        }
    }

    status = idap->GrowInternal(id_count);

    if (status == ZX_OK) {
        *ida_out = std::move(idap);
    }

    return status;
}

} // namespace id_allocator