/*
 * Copyright (c) 2024 Huawei Device Co., Ltd.
 * Licensed 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.
 */

#include "ecmascript/mem/heap-inl.h"
#include "ecmascript/mem/jit_fort.h"
#include "ecmascript/jit/jit.h"
#include "ecmascript/platform/backtrace.h"
#include "ecmascript/platform/file.h"
#include "ecmascript/platform/os.h"
#if defined(JIT_ENABLE_CODE_SIGN) && !defined(JIT_FORT_DISABLE)
#include <sys/ioctl.h>
#include <sys/prctl.h>

#define XPM_JITFORT_ENABLE_OPCODE 3
#define XPM_MAGIC 'x'
#define XPM_SET_JITFORT_ENABLE _IOW(XPM_MAGIC, XPM_JITFORT_ENABLE_OPCODE, unsigned long)
#endif

namespace panda::ecmascript {

template <>
FreeListAllocator<MemDesc>::FreeListAllocator(BaseHeap *heap, MemDescPool *pool, JitFort *fort)
    : memDescPool_(pool), heap_(heap)
{
    freeList_ = std::make_unique<FreeObjectList<MemDesc>>(fort);
}

JitFort::JitFort()
{
    jitFortMem_ = PageMap(JIT_FORT_REG_SPACE_MAX,
                          PageProtectProt(Jit::GetInstance()->IsDisableCodeSign() || !IsResourceAvailable()),
                          DEFAULT_REGION_SIZE, nullptr, MAP_JITFORT, true);
    jitFortBegin_ = reinterpret_cast<uintptr_t>(jitFortMem_.GetMem());
    jitFortSize_ = JIT_FORT_REG_SPACE_MAX;
    memDescPool_ = new MemDescPool(jitFortBegin_, jitFortSize_);
    allocator_ = new FreeListAllocator<MemDesc>(nullptr, memDescPool_, this);
    InitRegions();
    PrctlSetVMA(jitFortMem_.GetMem(), jitFortSize_, "ArkTS Code Jit");
    LOG_JIT(DEBUG) << "Small JitFort Begin " << (void *)jitFortBegin_ << " end "
                   << (void *)(jitFortBegin_ + jitFortSize_);
}

void JitFort::InitHugeRegion()
{
    if (hugeInUse_)
        return;
    hugeJitFortMem_ = PageMap(HUGE_JITFORT_REGION_SIZE,
                              PageProtectProt(Jit::GetInstance()->IsDisableCodeSign() || !IsResourceAvailable()),
                              HUGE_JITFORT_REGION_SIZE, nullptr, MAP_JITFORT, true);

    hugeJitFortBegin_ = reinterpret_cast<uintptr_t>(hugeJitFortMem_.GetMem());
    hugeJitFortSize_ = HUGE_JITFORT_REGION_SIZE;
    hugeMemDescPool_ = new MemDescPool(hugeJitFortBegin_, hugeJitFortSize_);
    hugeRegion_ = new JitFortRegion(nullptr, hugeJitFortBegin_, hugeJitFortBegin_ + HUGE_JITFORT_REGION_SIZE,
                                    RegionSpaceFlag::IN_HUGE_MACHINE_CODE_SPACE, hugeMemDescPool_);
    hugeAlloc_ = new FreeListAllocator<MemDesc>(nullptr, hugeMemDescPool_, this);
    hugeAlloc_->AddFree(hugeRegion_);
    hugeInUse_ = true;
    PrctlSetVMA(hugeJitFortMem_.GetMem(), HUGE_JITFORT_REGION_SIZE, "ArkTS Code Jit-Huge");
    LOG_JIT(DEBUG) << "Huge JitFort Begin " << (void *)hugeJitFortBegin_ << " end "
                   << (void *)(hugeJitFortBegin_ + hugeJitFortSize_);
}

JitFort::~JitFort()
{
    constexpr size_t numRegions = JIT_FORT_REG_SPACE_MAX / DEFAULT_REGION_SIZE;
    for (size_t i = 0; i < numRegions; i++) {
        if (regions_[i] != nullptr) {
            regions_[i]->DestroyFreeObjectSets();
            delete regions_[i];
        }
    }
    if (allocator_ != nullptr) {
        delete allocator_;
    }
    if (memDescPool_ != nullptr) {
        delete memDescPool_;
    }
    PageUnmap(jitFortMem_, true);
    if (hugeInUse_) {
        hugeRegion_->DestroyFreeObjectSets();
        delete hugeRegion_;
        delete hugeAlloc_;
        delete hugeMemDescPool_;
        PageUnmap(hugeJitFortMem_, true);
    }
}

void JitFort::InitRegions()
{
    for (size_t i = 0; i < MAX_JIT_FORT_REGIONS; i++) {
        uintptr_t mem = reinterpret_cast<uintptr_t>(jitFortMem_.GetMem()) + i*DEFAULT_REGION_SIZE;
        uintptr_t end = mem + DEFAULT_REGION_SIZE;
        JitFortRegion *region = new JitFortRegion(nullptr, mem, end, RegionSpaceFlag::IN_MACHINE_CODE_SPACE,
            memDescPool_);
        regions_[i] = region;
        LOG_JIT(DEBUG) << "Small JitFort Region " << i << std::hex << " init, 0x" << mem << " - 0x" << end
                       << ", size: 0x" << DEFAULT_REGION_SIZE << ", bitsetSize: " << region->GetGCBitsetSize();
    }
    AddRegion();
}

bool JitFort::AddRegion()
{
    if (nextFreeRegionIdx_ < MAX_JIT_FORT_REGIONS) {
        allocator_->AddFree(regions_[nextFreeRegionIdx_]);
        regionList_.AddNode(regions_[nextFreeRegionIdx_]);
        nextFreeRegionIdx_++;
        return true;
    }
    return false;
}

uintptr_t JitFort::AllocateSmall(size_t size)
{
    uintptr_t ret = ToUintPtr(nullptr);
    ret = allocator_->Allocate(size);
    if (ret == ToUintPtr(nullptr)) {
        if (AddRegion()) {
            ret = allocator_->Allocate(size);
        }
    }
    return ret;
}

uintptr_t JitFort::AllocateHuge(size_t size)
{
    if (UNLIKELY(!hugeInUse_)) {
        InitHugeRegion();
    }
    return hugeAlloc_->Allocate(size);
}

// Fort buf allocation is in multiples of FORT_BUF_ALIGN
size_t JitFort::FortAllocSize(size_t instrSize)
{
    return(AlignUp(instrSize, FORT_BUF_ALIGN));
}

uintptr_t JitFort::Allocate(MachineCodeDesc *desc)
{
    LockHolder lock(mutex_);

    size_t size = FortAllocSize(desc->instructionsSize);
    uintptr_t ret = ToUintPtr(nullptr);
    bool inHuge = desc->isHugeObj;
    if (UNLIKELY(inHuge)) {
        ret = AllocateHuge(size);
    } else {
        ret = AllocateSmall(size);
    }
    if (ret == ToUintPtr(nullptr)) {
        LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "JitFort:: Allocate return nullptr for size " << size;
        return ret;
    }
    // Record allocation to keep it from being collected by the next
    // JitFort::UpdateFreeSpace in case corresponding Machine code object is not
    // marked for sweep yet by then.
    MarkJitFortMemAwaitInstall(ret, size, inHuge);
    ASSERT((ret & FORT_BUF_ADDR_MASK) == 0);
    LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "JitFort:: Allocate " << (void *)ret << " - "
                   << (void *)(ret + size - 1) << " size " << size << " instructionsSize " << desc->instructionsSize;
    return ret;
}

// Called by GC thread during Mark to mark Fort buf alive.
// Encoding details in GCBitset for marking live Fort buffers:
//   Begin of a live Fort buf:
//     - 10: a live Jit Fort buf that has been installed
//     - 11: a live Jit Fort buf that has not been installed yet
//   End of a live Fort buf:
//     - 01
// This encoding requires 4 GCBitset bits to mark a live JitFort
// buffer, which makes the minimum Fort buffer allocation size 32 bytes.
void JitFort::MarkJitFortMemAlive(MachineCode *obj, bool inHuge)
{
    LockHolder lock(mutex_);
    size_t size = FortAllocSize(obj->GetInstructionsSize());
    uintptr_t addr = obj->GetText();
    uintptr_t endAddr = addr + size - 1;
    uint32_t regionIdx = -1;
    JitFortRegion *region = nullptr;
    if (LIKELY(!inHuge)) {
        regionIdx = AddrToFortRegionIdx(addr);
        region = regions_.at(regionIdx);
    } else {
        region = hugeRegion_;
    }
    region->AtomicMark(reinterpret_cast<void *>(addr));
    region->AtomicMark(reinterpret_cast<void *>(endAddr));

    LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "MarkFortMemAlive: addr: " << (void *)addr << " size: 0x"
                   << std::hex << size << " regionIdx: " << regionIdx << " instructionsSize: 0x"
                   << obj->GetInstructionsSize() << " mahine code addr: " << obj;
}

// Called by Jit Compile thread during JitFort Allocate to mark Fort buf
// awaiting install. Need mutex (JitGCLockHolder) for region gcBitset access.
// See JitFort::MarkJitFortMemAlive comments for mark bit encoding in GC bitset.
void JitFort::MarkJitFortMemAwaitInstall(uintptr_t addr, size_t size, bool inHuge)
{
    uintptr_t endAddr = addr + size - 1;
    uint32_t regionIdx = -1;
    JitFortRegion *region = nullptr;
    if (LIKELY(!inHuge)) {
        regionIdx = AddrToFortRegionIdx(addr);
        if (!InSmallRange(addr) || regionIdx >= MAX_JIT_FORT_REGIONS) {
            LOG_JIT(FATAL) << "JitFort::AddrToFortRegionIdx error, addr: 0x" << std::hex << addr
                           << ", jitFortBegin_: 0x" << jitFortBegin_ << ", regionIdx: " << regionIdx
                           << ", regions_.len: " << regions_.size();
        }
        region = regions_.at(regionIdx);
    } else {
        region = hugeRegion_;
    }
    region->AtomicMark(reinterpret_cast<void *>(addr));
    region->AtomicMark(reinterpret_cast<void *>(addr + sizeof(uint64_t))); // mark next bit
    region->AtomicMark(reinterpret_cast<void *>(endAddr));
    LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "MarkFortMemAwaitInstall: addr: " << (void *)addr << " size: 0x"
                   << std::hex << size << " regionIdx: " << regionIdx;
}

// Called by JS/Main thread during SafePoint to clear Fort buf AwaitInstall bit
// See JitFort::MarkJitFortMemAlive comments for mark bit encoding in GC bitset.
void JitFort::MarkJitFortMemInstalled(MachineCode *obj, bool inHuge)
{
    LockHolder lock(mutex_);
    size_t size = FortAllocSize(obj->GetInstructionsSize());
    uintptr_t addr = obj->GetText();
    uint32_t regionIdx = -1;
    JitFortRegion *region = nullptr;
    if (LIKELY(!inHuge)) {
        regionIdx = AddrToFortRegionIdx(addr);
        region = regions_.at(regionIdx);
    } else {
        region = hugeRegion_;
    }
    region->ClearMark(reinterpret_cast<void *>(addr + sizeof(uint64_t))); // clear next bit
    LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "MarkFortMemInstalled: addr: " << (void *)addr << " size: 0x"
                   << std::hex << size << " regionIdx: " << regionIdx << " instructionsSize: 0x"
                   << obj->GetInstructionsSize() << " mahine code addr: " << obj;
}

uint32_t JitFort::AddrToFortRegionIdx(uint64_t addr)
{
    ASSERT(InSmallRange(addr));
    uint32_t regionIdx = ((addr - jitFortBegin_) & ~(DEFAULT_REGION_MASK)) >> REGION_SIZE_LOG2;
    ASSERT(regionIdx < MAX_JIT_FORT_REGIONS);
    return regionIdx;
}

/*
 * Called from GC worker thread duing Old/Full GC Sweep (AsyncSweep). Mutex is needed
 * to ensure exclusive access to JitFort memory by GC thread when it frees JitFort mem
 * blocks, and by Jit compiled thread when it allocates Fort mem.
 */
void JitFort::UpdateFreeSpace(bool inHuge)
{
    if (!Jit::GetInstance()->IsEnableJitFort()) {
        return;
    }

    LockHolder lock(mutex_);

    // solve huge fort
    if (UNLIKELY(inHuge)) {
        if (LIKELY(!hugeInUse_)) {
            return;
        }
        hugeAlloc_->RebuildFreeList();
        FreeRegion(hugeRegion_, true);
    } else {
        if (!regionList_.GetLength()) {
            return;
        }
        allocator_->RebuildFreeList();
        JitFortRegion *region = regionList_.GetFirst();
        while (region) {
            FreeRegion(region);
            region = region->GetNext();
        }
    }
}

void JitFort::ClearMarkBits()
{
    ASSERT(g_isEnableCMCGC);
    LOG_JIT(INFO) << "JitFort::ClearMarkBits";
    JitFortRegion *region = regionList_.GetFirst();
    while (region) {
        region->IterateMarkedBitsConst(
            [](void *, size_t) {}); // Dummy vistor, only need to trigger bits clearing.
        region = region->GetNext();
    }
}

void JitFort::FreeRegion(JitFortRegion *region, bool inHuge)
{
    LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "JitFort::FreeRegion " << (void *)(region->GetBegin());

    uintptr_t freeStart = region->GetBegin();
    region->IterateMarkedBitsConst([this, &region, &freeStart, &inHuge](void *mem, size_t size) {
        ASSERT(region->InRange(ToUintPtr(mem)));
        (void)region;
        uintptr_t freeEnd = ToUintPtr(mem);
        if (freeStart != freeEnd) {
            LOG_JIT(INFO) << (inHuge ? "Huge " : "Small ")
                          << "JitFort::FreeRegion, freeStart: " << reinterpret_cast<void *>(freeStart)
                          << ", size: " << std::hex << freeEnd - freeStart;
            auto *allocator = inHuge ? hugeAlloc_ : allocator_;
            allocator->Free(freeStart, freeEnd - freeStart, true);
        }
        freeStart = freeEnd + size;
    });
    uintptr_t freeEnd = region->GetEnd();
    if (freeStart != freeEnd) {
        LOG_JIT(INFO) << (inHuge ? "Huge " : "Small ")
                      << "JitFort::FreeRegion last, freeStart: " << reinterpret_cast<void *>(freeStart)
                      << ", size: " << std::hex << freeEnd - freeStart;
        auto *allocator = inHuge ? hugeAlloc_ : allocator_;
        allocator->Free(freeStart, freeEnd - freeStart, true);
    }
}

bool JitFort::InJitFortRange(uintptr_t address) const
{
    return InSmallRange(address) || InHugeRange(address);
}

bool JitFort::InSmallRange(uintptr_t address) const
{
    return address >= jitFortBegin_ && (jitFortBegin_ + jitFortSize_) > 1 &&
        address <= (jitFortBegin_ + jitFortSize_ - 1);
}

bool JitFort::InHugeRange(uintptr_t address) const
{
    return address >= hugeJitFortBegin_ && (hugeJitFortBegin_ + hugeJitFortSize_) > 1 &&
        address <= (hugeJitFortBegin_ + hugeJitFortSize_ - 1);
}

void JitFort::PrepareSweeping()
{
    isSweeping_.store(false, std::memory_order_release);
}

// concurrent sweep - only one of the AsyncSweep task will do JitFort sweep
void JitFort::AsyncSweep(bool inHuge)
{
    bool expect = false;
    if (isSweeping_.compare_exchange_strong(expect, true, std::memory_order_seq_cst)) {
        LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "JitFort::AsyncSweep ";
        UpdateFreeSpace(inHuge);
    }
}

// non-concurrent sweep
void JitFort::Sweep(bool inHuge)
{
    LOG_JIT(DEBUG) << (inHuge ? "Huge " : "Small ") << "JitFort::Sweep";
    UpdateFreeSpace(inHuge);
}

// Used by JitFort::UpdateFreeSpace call path to find corresponding
// JitFortRegion for a free block in JitFort space, in order to put the blk into
// the corresponding free set of the JitFortRegion the free block belongs.
JitFortRegion *JitFort::ObjectAddressToRange(uintptr_t addr)
{
    ASSERT(InJitFortRange(addr));
    if (InSmallRange(addr)) {
        return regions_[AddrToFortRegionIdx(addr)];
    }
    return hugeRegion_;
}

template <size_t RegionMask>
void JitFortGCBitset<RegionMask>::MarkStartAddr(bool awaitInstall, uintptr_t startAddr, uint32_t index, uint32_t &word)
{
    if (!awaitInstall) {
        ClearMark(startAddr);
        word &= ~(1u << index);
    } else {
        word &= ~(1u << index);
        word &= ~(1u << (index+1));
        if (index + 1 == BIT_PER_WORD) {
            LOG_JIT(FATAL) << "JitFort::MarkStartAddr";
        }
    }
}

template <size_t RegionMask>
void JitFortGCBitset<RegionMask>::MarkEndAddr(bool awaitInstall, uintptr_t endAddr, uint32_t index, uint32_t &word)
{
    if (!awaitInstall) {
        ClearMark(endAddr - 1);
    }
    word &= ~(1u << index);
}

// See JitFort::MarkJitFortMemAlive comments for mark bit encoding in JitFort GC bitset.
template <size_t RegionMask>
template <typename Visitor>
void JitFortGCBitset<RegionMask>::IterateMarkedBitsConst(uintptr_t regionAddr, size_t bitsetSize, Visitor visitor)
{
    bool awaitInstall = false;
    uintptr_t startAddr = 0;
    uintptr_t endAddr = 0;

    auto words = Words();
    uint32_t wordCount = WordCount(bitsetSize);
    uint32_t index = BIT_PER_WORD;
    for (uint32_t i = 0; i < wordCount; i++) {
        uint32_t word = words[i];
        while (word != 0) {
            index = static_cast<uint32_t>(__builtin_ctz(word));
            if (index >= BIT_PER_WORD) {
                LOG_JIT(FATAL) << "Unexpected index from ctz: " << index << ", word: 0x" << std::hex << word;
            }
            LOG_JIT(DEBUG) << "i: " << i << ", index: " << index << ", word: 0x" << std::hex << word
                           << ", startAddr: 0x" << startAddr << ", endAddr: 0x" << endAddr;
            ASSERT(index < BIT_PER_WORD);
            if (!startAddr) {
                startAddr = regionAddr + (index << TAGGED_TYPE_SIZE_LOG);
                awaitInstall = Test(regionAddr + ((index+1) << TAGGED_TYPE_SIZE_LOG));
                MarkStartAddr(awaitInstall, startAddr, index, word);
            } else {
                endAddr = regionAddr + ((index+1) << TAGGED_TYPE_SIZE_LOG);
                LOG_JIT(DEBUG) << "Live Jit Mem " << (void *)startAddr << " size " << endAddr-startAddr;
                visitor(reinterpret_cast<void *>(startAddr), endAddr - startAddr);
                MarkEndAddr(awaitInstall, endAddr, index, word);
                awaitInstall = false;
                startAddr = 0;
            }
        }
        regionAddr += TAGGED_TYPE_SIZE * BIT_PER_WORD;
    }
}

bool JitFort::isResourceAvailable_ = true;
bool JitFort::IsResourceAvailable()
{
    return isResourceAvailable_;
}

void JitFort::InitJitFort()
{
#if defined(JIT_ENABLE_CODE_SIGN) && !defined(JIT_FORT_DISABLE)
    ECMA_BYTRACE_NAME(HITRACE_LEVEL_COMMERCIAL, HITRACE_TAG_ARK, "JIT::InitJitFortResource", "");
    int fd = open("/dev/xpm", O_RDWR);
    if (fd < 0) {
        isResourceAvailable_ = false;
        LOG_JIT(ERROR) << "Failed to init jitfort resource, open xpm failed: " << strerror(errno);
        return;
    }
    FdsanExchangeOwnerTag(reinterpret_cast<fd_t>(fd));
    int rc = ioctl(fd, XPM_SET_JITFORT_ENABLE, 0);
    if (rc < 0) {
        isResourceAvailable_ = false;
        LOG_JIT(ERROR) << "Failed to init jitfort resource, enable xpm failed: " << strerror(errno);
        Close(reinterpret_cast<fd_t>(fd));
        return;
    }
    Close(reinterpret_cast<fd_t>(fd));
#endif
}

void JitFort::InitJitFortResource()
{
#if defined(JIT_ENABLE_CODE_SIGN) && !defined(JIT_FORT_DISABLE)
    ECMA_BYTRACE_NAME(HITRACE_LEVEL_COMMERCIAL, HITRACE_TAG_ARK, "JIT::InitJitFortResource", "");
    constexpr int prSetJitFort = 0x6a6974;
    constexpr int jitFortInit = 5;
    int res = prctl(prSetJitFort, jitFortInit, 0);
    if (res < 0) {
        isResourceAvailable_ = false;
        LOG_JIT(ERROR) << "Failed to init jitfort resource: " << strerror(errno);
        return;
    }
    res = prctl(prSetJitFort, jitFortInit, 0);
    if (res >= 0 || errno != EEXIST) {
        isResourceAvailable_ = false;
        LOG_JIT(ERROR) << "jitfort not support";
    }
#endif
}

MemDescPool::MemDescPool(uintptr_t fortBegin, size_t fortSize)
    : fortBegin_(fortBegin), fortSize_(fortSize)
{
    Expand();
}

MemDescPool::~MemDescPool()
{
    for (const auto& block : memDescBlocks_) {
        if (block) {
            free(block);
        }
    }
}

MemDesc *MemDescPool::GetDesc()
{
    if (IsEmpty(freeList_)) {
        Expand();
    }
    if (!IsEmpty(freeList_)) { // LCOV_EXCL_BR_LINE
        MemDesc *res = freeList_;
        freeList_ = freeList_->GetNext();
        allocated_++;
        if (allocated_-returned_ > highwater_) {
            highwater_ = allocated_ - returned_;
        }
        return res;
    }
    return nullptr;
}

void MemDescPool::Expand()
{
    void *block = malloc(sizeof(MemDesc) * MEMDESCS_PER_BLOCK);
    if (block) {
        memDescBlocks_.push_back(block);
        for (size_t i = 0; i < MEMDESCS_PER_BLOCK; ++i) {
            Add(new (ToVoidPtr(reinterpret_cast<uintptr_t>(block) + i*sizeof(MemDesc))) MemDesc());
        }
    }
}

void MemDescPool::Add(MemDesc *desc)
{
    desc->SetNext(freeList_);
    freeList_ = desc;
}

}  // namespace panda::ecmascript