// 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.

#pragma once

#include <limits.h>
#include <stddef.h>
#include <stdint.h>

#include <zircon/process.h>
#include <zircon/types.h>
#include <fbl/algorithm.h>
#include <fbl/alloc_checker.h>
#include <fbl/array.h>
#include <fbl/macros.h>

#if !defined _KERNEL && defined __Fuchsia__
#include <lib/zx/vmo.h>
#include <zircon/syscalls.h>
#endif

namespace bitmap {

class DefaultStorage {
public:
    DISALLOW_COPY_ASSIGN_AND_MOVE(DefaultStorage);
    DefaultStorage() = default;

    zx_status_t Allocate(size_t size) {
        fbl::AllocChecker ac;
        auto arr = new (&ac) uint8_t[size];
        if (!ac.check()) {
            return ZX_ERR_NO_MEMORY;
        }
        storage_.reset(arr, size);
        return ZX_OK;
    }
    void* GetData() { return storage_.get(); }
    const void* GetData() const { return storage_.get(); }
private:
    fbl::Array<uint8_t> storage_;
};

template <size_t N>
class FixedStorage {
public:
    DISALLOW_COPY_ASSIGN_AND_MOVE(FixedStorage);
    FixedStorage() = default;

    zx_status_t Allocate(size_t size) {
        ZX_ASSERT(size <= N);
        return ZX_OK;
    }
    void* GetData() { return storage_; }
    const void* GetData() const { return storage_; }
private:
    size_t storage_[(N + sizeof(size_t) - 1) / sizeof(size_t)];
};

#if !defined _KERNEL && defined __Fuchsia__
class VmoStorage {
public:
    DISALLOW_COPY_AND_ASSIGN_ALLOW_MOVE(VmoStorage);
    VmoStorage() :
        vmo_(ZX_HANDLE_INVALID),
        mapped_addr_(0),
        size_(0) {}

    VmoStorage(VmoStorage&& rhs)
        : vmo_(std::move(rhs.vmo_)), mapped_addr_(rhs.mapped_addr_), size_(rhs.size_) {
        rhs.mapped_addr_ = 0;
    }

    VmoStorage& operator=(VmoStorage&& rhs) {
        Release();
        vmo_ = std::move(rhs.vmo_);
        mapped_addr_ = rhs.mapped_addr_;
        size_ = rhs.size_;
        rhs.mapped_addr_ = 0;
        return *this;
    }

    ~VmoStorage() {
        Release();
    }

    zx_status_t Allocate(size_t size) {
        Release();
        size_ = fbl::round_up(size, static_cast<size_t>(PAGE_SIZE));
        zx_status_t status;
        if ((status = zx::vmo::create(size_, ZX_VMO_RESIZABLE, &vmo_)) != ZX_OK) {
            return status;
        } else if ((status = zx_vmar_map(zx_vmar_root_self(),
                                         ZX_VM_PERM_READ | ZX_VM_PERM_WRITE,
                                         0, vmo_.get(), 0, size_, &mapped_addr_)) != ZX_OK) {
            vmo_.reset();
            return status;
        }
        vmo_.set_property(ZX_PROP_NAME, "vmo-backed-bitmap", strlen("vmo-backed-bitmap"));
        return ZX_OK;
    }

    zx_status_t Grow(size_t size) {
        if (size <= size_) {
            return ZX_OK;
        }

        size = fbl::round_up(size, static_cast<size_t>(PAGE_SIZE));
        zx_status_t status;
        if ((status = vmo_.set_size(size)) != ZX_OK) {
            return status;
        }


        zx_info_vmar_t vmar_info;
        if ((status = zx_object_get_info(zx_vmar_root_self(), ZX_INFO_VMAR,
                                         &vmar_info, sizeof(vmar_info), NULL,
                                         NULL)) != ZX_OK) {
            return status;
        }

        // Try to extend mapping
        uintptr_t addr;
        if ((status = zx_vmar_map(zx_vmar_root_self(),
                                  ZX_VM_PERM_READ | ZX_VM_PERM_WRITE |
                                  ZX_VM_SPECIFIC,
                                  mapped_addr_ + size_ -
                                  vmar_info.base, vmo_.get(), size_, size - size_,
                                  &addr)) != ZX_OK) {
            // If extension fails, create entirely new mapping and unmap the old one
            if ((status = zx_vmar_map(zx_vmar_root_self(),
                                      ZX_VM_PERM_READ | ZX_VM_PERM_WRITE,
                                      0, vmo_.get(), 0, size, &addr)) != ZX_OK) {
                return status;
            }

            if ((status = zx_vmar_unmap(zx_vmar_root_self(), mapped_addr_, size_)) != ZX_OK) {
                return status;
            }

            mapped_addr_ = addr;
        }

        return ZX_OK;
    }

    void* GetData() { ZX_DEBUG_ASSERT(mapped_addr_ != 0); return (void*) mapped_addr_; }
    const void* GetData() const { ZX_DEBUG_ASSERT(mapped_addr_ != 0); return (void*) mapped_addr_; }
    const zx::vmo& GetVmo() const { ZX_DEBUG_ASSERT(mapped_addr_ != 0); return vmo_; }
private:
    void Release() {
        if (mapped_addr_ != 0) {
            zx_vmar_unmap(zx_vmar_root_self(), mapped_addr_, size_);
        }
        mapped_addr_ = 0;
    }
    zx::vmo vmo_;
    uintptr_t mapped_addr_;
    size_t size_;
};
#endif

} // namespace bitmap