code:
// SPDX-FileCopyrightText: Copyright 2020 yuzu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#pragma once

#include "common/alignment.h"
#include "common/assert.h"
#include "common/intrusive_red_black_tree.h"
#include "core/hle/kernel/k_typed_address.h"
#include "core/hle/kernel/memory_types.h"
#include "core/hle/kernel/svc_types.h"

namespace Kernel {

enum class KMemoryState : u32 {
    None = 0,
    Mask = 0xFF,
    All = ~None,

    FlagCanReprotect = (1 << 8),
    FlagCanDebug = (1 << 9),
    FlagCanUseIpc = (1 << 10),
    FlagCanUseNonDeviceIpc = (1 << 11),
    FlagCanUseNonSecureIpc = (1 << 12),
    FlagMapped = (1 << 13),
    FlagCode = (1 << 14),
    FlagCanAlias = (1 << 15),
    FlagCanCodeAlias = (1 << 16),
    FlagCanTransfer = (1 << 17),
    FlagCanQueryPhysical = (1 << 18),
    FlagCanDeviceMap = (1 << 19),
    FlagCanAlignedDeviceMap = (1 << 20),
    FlagCanIpcUserBuffer = (1 << 21),
    FlagReferenceCounted = (1 << 22),
    FlagCanMapProcess = (1 << 23),
    FlagCanChangeAttribute = (1 << 24),
    FlagCanCodeMemory = (1 << 25),
    FlagLinearMapped = (1 << 26),
    FlagCanPermissionLock = (1 << 27),

    FlagsData = FlagCanReprotect | FlagCanUseIpc | FlagCanUseNonDeviceIpc | FlagCanUseNonSecureIpc |
                FlagMapped | FlagCanAlias | FlagCanTransfer | FlagCanQueryPhysical |
                FlagCanDeviceMap | FlagCanAlignedDeviceMap | FlagCanIpcUserBuffer |
                FlagReferenceCounted | FlagCanChangeAttribute | FlagLinearMapped,

    FlagsCode = FlagCanDebug | FlagCanUseIpc | FlagCanUseNonDeviceIpc | FlagCanUseNonSecureIpc |
                FlagMapped | FlagCode | FlagCanQueryPhysical | FlagCanDeviceMap |
                FlagCanAlignedDeviceMap | FlagReferenceCounted | FlagLinearMapped,

    FlagsMisc = FlagMapped | FlagReferenceCounted | FlagCanQueryPhysical | FlagCanDeviceMap |
                FlagLinearMapped,

    Free = static_cast<u32>(Svc::MemoryState::Free),

    IoMemory = static_cast<u32>(Svc::MemoryState::Io) | FlagMapped | FlagCanDeviceMap |
               FlagCanAlignedDeviceMap,
    IoRegister =
        static_cast<u32>(Svc::MemoryState::Io) | FlagCanDeviceMap | FlagCanAlignedDeviceMap,

    Static = static_cast<u32>(Svc::MemoryState::Static) | FlagMapped | FlagCanQueryPhysical,
    Code = static_cast<u32>(Svc::MemoryState::Code) | FlagsCode | FlagCanMapProcess,
    CodeData = static_cast<u32>(Svc::MemoryState::CodeData) | FlagsData | FlagCanMapProcess |
               FlagCanCodeMemory | FlagCanPermissionLock,
    Normal = static_cast<u32>(Svc::MemoryState::Normal) | FlagsData | FlagCanCodeMemory,
    Shared = static_cast<u32>(Svc::MemoryState::Shared) | FlagMapped | FlagReferenceCounted |
             FlagLinearMapped,

    // Alias was removed after 1.0.0.

    AliasCode = static_cast<u32>(Svc::MemoryState::AliasCode) | FlagsCode | FlagCanMapProcess |
                FlagCanCodeAlias,
    AliasCodeData = static_cast<u32>(Svc::MemoryState::AliasCodeData) | FlagsData |
                    FlagCanMapProcess | FlagCanCodeAlias | FlagCanCodeMemory |
                    FlagCanPermissionLock,

    Ipc = static_cast<u32>(Svc::MemoryState::Ipc) | FlagsMisc | FlagCanAlignedDeviceMap |
          FlagCanUseIpc | FlagCanUseNonSecureIpc | FlagCanUseNonDeviceIpc,

    Stack = static_cast<u32>(Svc::MemoryState::Stack) | FlagsMisc | FlagCanAlignedDeviceMap |
            FlagCanUseIpc | FlagCanUseNonSecureIpc | FlagCanUseNonDeviceIpc,

    ThreadLocal = static_cast<u32>(Svc::MemoryState::ThreadLocal) | FlagLinearMapped,

    Transferred = static_cast<u32>(Svc::MemoryState::Transferred) | FlagsMisc |
                  FlagCanAlignedDeviceMap | FlagCanChangeAttribute | FlagCanUseIpc |
                  FlagCanUseNonSecureIpc | FlagCanUseNonDeviceIpc,

    SharedTransferred = static_cast<u32>(Svc::MemoryState::SharedTransferred) | FlagsMisc |
                        FlagCanAlignedDeviceMap | FlagCanUseNonSecureIpc | FlagCanUseNonDeviceIpc,

    SharedCode = static_cast<u32>(Svc::MemoryState::SharedCode) | FlagMapped |
                 FlagReferenceCounted | FlagLinearMapped | FlagCanUseNonSecureIpc |
                 FlagCanUseNonDeviceIpc,

    Inaccessible = static_cast<u32>(Svc::MemoryState::Inaccessible),

    NonSecureIpc = static_cast<u32>(Svc::MemoryState::NonSecureIpc) | FlagsMisc |
                   FlagCanAlignedDeviceMap | FlagCanUseNonSecureIpc | FlagCanUseNonDeviceIpc,

    NonDeviceIpc =
        static_cast<u32>(Svc::MemoryState::NonDeviceIpc) | FlagsMisc | FlagCanUseNonDeviceIpc,

    Kernel = static_cast<u32>(Svc::MemoryState::Kernel),

    GeneratedCode = static_cast<u32>(Svc::MemoryState::GeneratedCode) | FlagMapped |
                    FlagReferenceCounted | FlagCanDebug | FlagLinearMapped,
    CodeOut = static_cast<u32>(Svc::MemoryState::CodeOut) | FlagMapped | FlagReferenceCounted |
              FlagLinearMapped,

    Coverage = static_cast<u32>(Svc::MemoryState::Coverage) | FlagMapped,

    Insecure = static_cast<u32>(Svc::MemoryState::Insecure) | FlagMapped | FlagReferenceCounted |
               FlagLinearMapped | FlagCanChangeAttribute | FlagCanDeviceMap |
               FlagCanAlignedDeviceMap | FlagCanQueryPhysical | FlagCanUseNonSecureIpc |
               FlagCanUseNonDeviceIpc,
};
DECLARE_ENUM_FLAG_OPERATORS(KMemoryState);

static_assert(static_cast<u32>(KMemoryState::Free) == 0x00000000);
static_assert(static_cast<u32>(KMemoryState::IoMemory) == 0x00182001);
static_assert(static_cast<u32>(KMemoryState::IoRegister) == 0x00180001);
static_assert(static_cast<u32>(KMemoryState::Static) == 0x00042002);
static_assert(static_cast<u32>(KMemoryState::Code) == 0x04DC7E03);
static_assert(static_cast<u32>(KMemoryState::CodeData) == 0x0FFEBD04);
static_assert(static_cast<u32>(KMemoryState::Normal) == 0x077EBD05);
static_assert(static_cast<u32>(KMemoryState::Shared) == 0x04402006);

static_assert(static_cast<u32>(KMemoryState::AliasCode) == 0x04DD7E08);
static_assert(static_cast<u32>(KMemoryState::AliasCodeData) == 0x0FFFBD09);
static_assert(static_cast<u32>(KMemoryState::Ipc) == 0x045C3C0A);
static_assert(static_cast<u32>(KMemoryState::Stack) == 0x045C3C0B);
static_assert(static_cast<u32>(KMemoryState::ThreadLocal) == 0x0400000C);
static_assert(static_cast<u32>(KMemoryState::Transferred) == 0x055C3C0D);
static_assert(static_cast<u32>(KMemoryState::SharedTransferred) == 0x045C380E);
static_assert(static_cast<u32>(KMemoryState::SharedCode) == 0x0440380F);
static_assert(static_cast<u32>(KMemoryState::Inaccessible) == 0x00000010);
static_assert(static_cast<u32>(KMemoryState::NonSecureIpc) == 0x045C3811);
static_assert(static_cast<u32>(KMemoryState::NonDeviceIpc) == 0x044C2812);
static_assert(static_cast<u32>(KMemoryState::Kernel) == 0x00000013);
static_assert(static_cast<u32>(KMemoryState::GeneratedCode) == 0x04402214);
static_assert(static_cast<u32>(KMemoryState::CodeOut) == 0x04402015);
static_assert(static_cast<u32>(KMemoryState::Coverage) == 0x00002016);
static_assert(static_cast<u32>(KMemoryState::Insecure) == 0x055C3817);

enum class KMemoryPermission : u8 {
    None = 0,
    All = static_cast<u8>(~None),

    KernelShift = 3,

    KernelRead = static_cast<u8>(Svc::MemoryPermission::Read) << KernelShift,
    KernelWrite = static_cast<u8>(Svc::MemoryPermission::Write) << KernelShift,
    KernelExecute = static_cast<u8>(Svc::MemoryPermission::Execute) << KernelShift,

    NotMapped = (1 << (2 * KernelShift)),

    KernelReadWrite = KernelRead | KernelWrite,
    KernelReadExecute = KernelRead | KernelExecute,

    UserRead = static_cast<u8>(Svc::MemoryPermission::Read) | KernelRead,
    UserWrite = static_cast<u8>(Svc::MemoryPermission::Write) | KernelWrite,
    UserExecute = static_cast<u8>(Svc::MemoryPermission::Execute),

    UserReadWrite = UserRead | UserWrite,
    UserReadExecute = UserRead | UserExecute,

    UserMask = static_cast<u8>(Svc::MemoryPermission::Read | Svc::MemoryPermission::Write |
                               Svc::MemoryPermission::Execute),

    IpcLockChangeMask = NotMapped | UserReadWrite,
};
DECLARE_ENUM_FLAG_OPERATORS(KMemoryPermission);

constexpr KMemoryPermission ConvertToKMemoryPermission(Svc::MemoryPermission perm) {
    return static_cast<KMemoryPermission>(
        (static_cast<KMemoryPermission>(perm) & KMemoryPermission::UserMask) |
        KMemoryPermission::KernelRead |
        ((static_cast<KMemoryPermission>(perm) & KMemoryPermission::UserWrite)
         << KMemoryPermission::KernelShift) |
        (perm == Svc::MemoryPermission::None ? KMemoryPermission::NotMapped
                                             : KMemoryPermission::None));
}

enum class KMemoryAttribute : u8 {
    None = 0x00,
    All = 0xFF,
    UserMask = All,

    Locked = static_cast<u8>(Svc::MemoryAttribute::Locked),
    IpcLocked = static_cast<u8>(Svc::MemoryAttribute::IpcLocked),
    DeviceShared = static_cast<u8>(Svc::MemoryAttribute::DeviceShared),
    Uncached = static_cast<u8>(Svc::MemoryAttribute::Uncached),
    PermissionLocked = static_cast<u8>(Svc::MemoryAttribute::PermissionLocked),

    SetMask = Uncached | PermissionLocked,
};
DECLARE_ENUM_FLAG_OPERATORS(KMemoryAttribute);

enum class KMemoryBlockDisableMergeAttribute : u8 {
    None = 0,
    Normal = (1u << 0),
    DeviceLeft = (1u << 1),
    IpcLeft = (1u << 2),
    Locked = (1u << 3),
    DeviceRight = (1u << 4),

    AllLeft = Normal | DeviceLeft | IpcLeft | Locked,
    AllRight = DeviceRight,
};
DECLARE_ENUM_FLAG_OPERATORS(KMemoryBlockDisableMergeAttribute);

struct KMemoryInfo {
    uintptr_t m_address;
    size_t m_size;
    KMemoryState m_state;
    u16 m_device_disable_merge_left_count;
    u16 m_device_disable_merge_right_count;
    u16 m_ipc_lock_count;
    u16 m_device_use_count;
    u16 m_ipc_disable_merge_count;
    KMemoryPermission m_permission;
    KMemoryAttribute m_attribute;
    KMemoryPermission m_original_permission;
    KMemoryBlockDisableMergeAttribute m_disable_merge_attribute;

    constexpr Svc::MemoryInfo GetSvcMemoryInfo() const {
        return {
            .base_address = m_address,
            .size = m_size,
            .state = static_cast<Svc::MemoryState>(m_state & KMemoryState::Mask),
            .attribute =
                static_cast<Svc::MemoryAttribute>(m_attribute & KMemoryAttribute::UserMask),
            .permission =
                static_cast<Svc::MemoryPermission>(m_permission & KMemoryPermission::UserMask),
            .ipc_count = m_ipc_lock_count,
            .device_count = m_device_use_count,
            .padding = {},
        };
    }

    constexpr uintptr_t GetAddress() const {
        return m_address;
    }

    constexpr size_t GetSize() const {
        return m_size;
    }

    constexpr size_t GetNumPages() const {
        return this->GetSize() / PageSize;
    }

    constexpr uintptr_t GetEndAddress() const {
        return this->GetAddress() + this->GetSize();
    }

    constexpr uintptr_t GetLastAddress() const {
        return this->GetEndAddress() - 1;
    }

    constexpr u16 GetIpcLockCount() const {
        return m_ipc_lock_count;
    }

    constexpr u16 GetIpcDisableMergeCount() const {
        return m_ipc_disable_merge_count;
    }

    constexpr KMemoryState GetState() const {
        return m_state;
    }

    constexpr Svc::MemoryState GetSvcState() const {
        return static_cast<Svc::MemoryState>(m_state & KMemoryState::Mask);
    }

    constexpr KMemoryPermission GetPermission() const {
        return m_permission;
    }

    constexpr KMemoryPermission GetOriginalPermission() const {
        return m_original_permission;
    }

    constexpr KMemoryAttribute GetAttribute() const {
        return m_attribute;
    }

    constexpr KMemoryBlockDisableMergeAttribute GetDisableMergeAttribute() const {
        return m_disable_merge_attribute;
    }
};

class KMemoryBlock : public Common::IntrusiveRedBlackTreeBaseNode<KMemoryBlock> {
private:
    u16 m_device_disable_merge_left_count{};
    u16 m_device_disable_merge_right_count{};
    KProcessAddress m_address{};
    size_t m_num_pages{};
    KMemoryState m_memory_state{KMemoryState::None};
    u16 m_ipc_lock_count{};
    u16 m_device_use_count{};
    u16 m_ipc_disable_merge_count{};
    KMemoryPermission m_permission{KMemoryPermission::None};
    KMemoryPermission m_original_permission{KMemoryPermission::None};
    KMemoryAttribute m_attribute{KMemoryAttribute::None};
    KMemoryBlockDisableMergeAttribute m_disable_merge_attribute{
        KMemoryBlockDisableMergeAttribute::None};

public:
    static constexpr int Compare(const KMemoryBlock& lhs, const KMemoryBlock& rhs) {
        if (lhs.GetAddress() < rhs.GetAddress()) {
            return -1;
        } else if (lhs.GetAddress() <= rhs.GetLastAddress()) {
            return 0;
        } else {
            return 1;
        }
    }

public:
    constexpr KProcessAddress GetAddress() const {
        return m_address;
    }

    constexpr size_t GetNumPages() const {
        return m_num_pages;
    }

    constexpr size_t GetSize() const {
        return this->GetNumPages() * PageSize;
    }

    constexpr KProcessAddress GetEndAddress() const {
        return this->GetAddress() + this->GetSize();
    }

    constexpr KProcessAddress GetLastAddress() const {
        return this->GetEndAddress() - 1;
    }

    constexpr KMemoryState GetState() const {
        return m_memory_state;
    }

    constexpr u16 GetIpcLockCount() const {
        return m_ipc_lock_count;
    }

    constexpr u16 GetIpcDisableMergeCount() const {
        return m_ipc_disable_merge_count;
    }

    constexpr KMemoryPermission GetPermission() const {
        return m_permission;
    }

    constexpr KMemoryPermission GetOriginalPermission() const {
        return m_original_permission;
    }

    constexpr KMemoryAttribute GetAttribute() const {
        return m_attribute;
    }

    constexpr KMemoryInfo GetMemoryInfo() const {
        return {
            .m_address = GetInteger(this->GetAddress()),
            .m_size = this->GetSize(),
            .m_state = m_memory_state,
            .m_device_disable_merge_left_count = m_device_disable_merge_left_count,
            .m_device_disable_merge_right_count = m_device_disable_merge_right_count,
            .m_ipc_lock_count = m_ipc_lock_count,
            .m_device_use_count = m_device_use_count,
            .m_ipc_disable_merge_count = m_ipc_disable_merge_count,
            .m_permission = m_permission,
            .m_attribute = m_attribute,
            .m_original_permission = m_original_permission,
            .m_disable_merge_attribute = m_disable_merge_attribute,
        };
    }

public:
    explicit KMemoryBlock() = default;

    constexpr KMemoryBlock(KProcessAddress addr, size_t np, KMemoryState ms, KMemoryPermission p,
                           KMemoryAttribute attr)
        : Common::IntrusiveRedBlackTreeBaseNode<KMemoryBlock>(), m_address(addr), m_num_pages(np),
          m_memory_state(ms), m_permission(p), m_attribute(attr) {}

    constexpr void Initialize(KProcessAddress addr, size_t np, KMemoryState ms, KMemoryPermission p,
                              KMemoryAttribute attr) {
        m_device_disable_merge_left_count = 0;
        m_device_disable_merge_right_count = 0;
        m_address = addr;
        m_num_pages = np;
        m_memory_state = ms;
        m_ipc_lock_count = 0;
        m_device_use_count = 0;
        m_permission = p;
        m_original_permission = KMemoryPermission::None;
        m_attribute = attr;
        m_disable_merge_attribute = KMemoryBlockDisableMergeAttribute::None;
    }

    constexpr bool HasProperties(KMemoryState s, KMemoryPermission p, KMemoryAttribute a) const {
        constexpr auto AttributeIgnoreMask =
            KMemoryAttribute::IpcLocked | KMemoryAttribute::DeviceShared;
        return m_memory_state == s && m_permission == p &&
               (m_attribute | AttributeIgnoreMask) == (a | AttributeIgnoreMask);
    }

    constexpr bool HasSameProperties(const KMemoryBlock& rhs) const {
        return m_memory_state == rhs.m_memory_state && m_permission == rhs.m_permission &&
               m_original_permission == rhs.m_original_permission &&
               m_attribute == rhs.m_attribute && m_ipc_lock_count == rhs.m_ipc_lock_count &&
               m_device_use_count == rhs.m_device_use_count;
    }

    constexpr bool CanMergeWith(const KMemoryBlock& rhs) const {
        return this->HasSameProperties(rhs) &&
               (m_disable_merge_attribute & KMemoryBlockDisableMergeAttribute::AllRight) ==
                   KMemoryBlockDisableMergeAttribute::None &&
               (rhs.m_disable_merge_attribute & KMemoryBlockDisableMergeAttribute::AllLeft) ==
                   KMemoryBlockDisableMergeAttribute::None;
    }

    constexpr bool Contains(KProcessAddress addr) const {
        return this->GetAddress() <= addr && addr <= this->GetEndAddress();
    }

    constexpr void Add(const KMemoryBlock& added_block) {
        ASSERT(added_block.GetNumPages() > 0);
        ASSERT(this->GetAddress() + added_block.GetSize() - 1 <
               this->GetEndAddress() + added_block.GetSize() - 1);

        m_num_pages += added_block.GetNumPages();
        m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
            m_disable_merge_attribute | added_block.m_disable_merge_attribute);
        m_device_disable_merge_right_count = added_block.m_device_disable_merge_right_count;
    }

    constexpr void Update(KMemoryState s, KMemoryPermission p, KMemoryAttribute a,
                          bool set_disable_merge_attr, u8 set_mask, u8 clear_mask) {
        ASSERT(m_original_permission == KMemoryPermission::None);
        ASSERT((m_attribute & KMemoryAttribute::IpcLocked) == KMemoryAttribute::None);

        m_memory_state = s;
        m_permission = p;
        m_attribute = static_cast<KMemoryAttribute>(
            a | (m_attribute & (KMemoryAttribute::IpcLocked | KMemoryAttribute::DeviceShared)));

        if (set_disable_merge_attr && set_mask != 0) {
            m_disable_merge_attribute = m_disable_merge_attribute |
                                        static_cast<KMemoryBlockDisableMergeAttribute>(set_mask);
        }
        if (clear_mask != 0) {
            m_disable_merge_attribute = m_disable_merge_attribute &
                                        static_cast<KMemoryBlockDisableMergeAttribute>(~clear_mask);
        }
    }

    constexpr void UpdateAttribute(KMemoryAttribute mask, KMemoryAttribute attr) {
        ASSERT(False(mask & KMemoryAttribute::IpcLocked));
        ASSERT(False(mask & KMemoryAttribute::DeviceShared));

        m_attribute = (m_attribute & ~mask) | attr;
    }

    constexpr void Split(KMemoryBlock* block, KProcessAddress addr) {
        ASSERT(this->GetAddress() < addr);
        ASSERT(this->Contains(addr));
        ASSERT(Common::IsAligned(GetInteger(addr), PageSize));

        block->m_address = m_address;
        block->m_num_pages = (addr - this->GetAddress()) / PageSize;
        block->m_memory_state = m_memory_state;
        block->m_ipc_lock_count = m_ipc_lock_count;
        block->m_device_use_count = m_device_use_count;
        block->m_permission = m_permission;
        block->m_original_permission = m_original_permission;
        block->m_attribute = m_attribute;
        block->m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
            m_disable_merge_attribute & KMemoryBlockDisableMergeAttribute::AllLeft);
        block->m_ipc_disable_merge_count = m_ipc_disable_merge_count;
        block->m_device_disable_merge_left_count = m_device_disable_merge_left_count;
        block->m_device_disable_merge_right_count = 0;

        m_address = addr;
        m_num_pages -= block->m_num_pages;

        m_ipc_disable_merge_count = 0;
        m_device_disable_merge_left_count = 0;
        m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
            m_disable_merge_attribute & KMemoryBlockDisableMergeAttribute::AllRight);
    }

    constexpr void UpdateDeviceDisableMergeStateForShareLeft(KMemoryPermission new_perm, bool left,
                                                             bool right) {
        // New permission/right aren't used.
        if (left) {
            m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
                m_disable_merge_attribute | KMemoryBlockDisableMergeAttribute::DeviceLeft);
            const u16 new_device_disable_merge_left_count = ++m_device_disable_merge_left_count;
            ASSERT(new_device_disable_merge_left_count > 0);
        }
    }

    constexpr void UpdateDeviceDisableMergeStateForShareRight(KMemoryPermission new_perm, bool left,
                                                              bool right) {
        // New permission/left aren't used.
        if (right) {
            m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
                m_disable_merge_attribute | KMemoryBlockDisableMergeAttribute::DeviceRight);
            const u16 new_device_disable_merge_right_count = ++m_device_disable_merge_right_count;
            ASSERT(new_device_disable_merge_right_count > 0);
        }
    }

    constexpr void UpdateDeviceDisableMergeStateForShare(KMemoryPermission new_perm, bool left,
                                                         bool right) {
        this->UpdateDeviceDisableMergeStateForShareLeft(new_perm, left, right);
        this->UpdateDeviceDisableMergeStateForShareRight(new_perm, left, right);
    }

    constexpr void ShareToDevice(KMemoryPermission new_perm, bool left, bool right) {
        // New permission isn't used.

        // We must either be shared or have a zero lock count.
        ASSERT((m_attribute & KMemoryAttribute::DeviceShared) == KMemoryAttribute::DeviceShared ||
               m_device_use_count == 0);

        // Share.
        const u16 new_count = ++m_device_use_count;
        ASSERT(new_count > 0);

        m_attribute = static_cast<KMemoryAttribute>(m_attribute | KMemoryAttribute::DeviceShared);

        this->UpdateDeviceDisableMergeStateForShare(new_perm, left, right);
    }

    constexpr void UpdateDeviceDisableMergeStateForUnshareLeft(KMemoryPermission new_perm,
                                                               bool left, bool right) {
        // New permission/right aren't used.

        if (left) {
            if (!m_device_disable_merge_left_count) {
                return;
            }
            --m_device_disable_merge_left_count;
        }

        m_device_disable_merge_left_count =
            std::min(m_device_disable_merge_left_count, m_device_use_count);

        if (m_device_disable_merge_left_count == 0) {
            m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
                m_disable_merge_attribute & ~KMemoryBlockDisableMergeAttribute::DeviceLeft);
        }
    }

    constexpr void UpdateDeviceDisableMergeStateForUnshareRight(KMemoryPermission new_perm,
                                                                bool left, bool right) {
        // New permission/left aren't used.

        if (right) {
            const u16 old_device_disable_merge_right_count = m_device_disable_merge_right_count--;
            ASSERT(old_device_disable_merge_right_count > 0);
            if (old_device_disable_merge_right_count == 1) {
                m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
                    m_disable_merge_attribute & ~KMemoryBlockDisableMergeAttribute::DeviceRight);
            }
        }
    }

    constexpr void UpdateDeviceDisableMergeStateForUnshare(KMemoryPermission new_perm, bool left,
                                                           bool right) {
        this->UpdateDeviceDisableMergeStateForUnshareLeft(new_perm, left, right);
        this->UpdateDeviceDisableMergeStateForUnshareRight(new_perm, left, right);
    }

    constexpr void UnshareToDevice(KMemoryPermission new_perm, bool left, bool right) {
        // New permission isn't used.

        // We must be shared.
        ASSERT((m_attribute & KMemoryAttribute::DeviceShared) == KMemoryAttribute::DeviceShared);

        // Unhare.
        const u16 old_count = m_device_use_count--;
        ASSERT(old_count > 0);

        if (old_count == 1) {
            m_attribute =
                static_cast<KMemoryAttribute>(m_attribute & ~KMemoryAttribute::DeviceShared);
        }

        this->UpdateDeviceDisableMergeStateForUnshare(new_perm, left, right);
    }

    constexpr void UnshareToDeviceRight(KMemoryPermission new_perm, bool left, bool right) {
        // New permission isn't used.

        // We must be shared.
        ASSERT((m_attribute & KMemoryAttribute::DeviceShared) == KMemoryAttribute::DeviceShared);

        // Unhare.
        const u16 old_count = m_device_use_count--;
        ASSERT(old_count > 0);

        if (old_count == 1) {
            m_attribute =
                static_cast<KMemoryAttribute>(m_attribute & ~KMemoryAttribute::DeviceShared);
        }

        this->UpdateDeviceDisableMergeStateForUnshareRight(new_perm, left, right);
    }

    constexpr void LockForIpc(KMemoryPermission new_perm, bool left, bool right) {
        // We must either be locked or have a zero lock count.
        ASSERT((m_attribute & KMemoryAttribute::IpcLocked) == KMemoryAttribute::IpcLocked ||
               m_ipc_lock_count == 0);

        // Lock.
        const u16 new_lock_count = ++m_ipc_lock_count;
        ASSERT(new_lock_count > 0);

        // If this is our first lock, update our permissions.
        if (new_lock_count == 1) {
            ASSERT(m_original_permission == KMemoryPermission::None);
            ASSERT((m_permission | new_perm | KMemoryPermission::NotMapped) ==
                   (m_permission | KMemoryPermission::NotMapped));
            ASSERT((m_permission & KMemoryPermission::UserExecute) !=
                       KMemoryPermission::UserExecute ||
                   (new_perm == KMemoryPermission::UserRead));
            m_original_permission = m_permission;
            m_permission = static_cast<KMemoryPermission>(
                (new_perm & KMemoryPermission::IpcLockChangeMask) |
                (m_original_permission & ~KMemoryPermission::IpcLockChangeMask));
        }
        m_attribute = static_cast<KMemoryAttribute>(m_attribute | KMemoryAttribute::IpcLocked);

        if (left) {
            m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
                m_disable_merge_attribute | KMemoryBlockDisableMergeAttribute::IpcLeft);
            const u16 new_ipc_disable_merge_count = ++m_ipc_disable_merge_count;
            ASSERT(new_ipc_disable_merge_count > 0);
        }
    }

    constexpr void UnlockForIpc(KMemoryPermission new_perm, bool left, bool right) {
        // New permission isn't used.

        // We must be locked.
        ASSERT((m_attribute & KMemoryAttribute::IpcLocked) == KMemoryAttribute::IpcLocked);

        // Unlock.
        const u16 old_lock_count = m_ipc_lock_count--;
        ASSERT(old_lock_count > 0);

        // If this is our last unlock, update our permissions.
        if (old_lock_count == 1) {
            ASSERT(m_original_permission != KMemoryPermission::None);
            m_permission = m_original_permission;
            m_original_permission = KMemoryPermission::None;
            m_attribute = static_cast<KMemoryAttribute>(m_attribute & ~KMemoryAttribute::IpcLocked);
        }

        if (left) {
            const u16 old_ipc_disable_merge_count = m_ipc_disable_merge_count--;
            ASSERT(old_ipc_disable_merge_count > 0);
            if (old_ipc_disable_merge_count == 1) {
                m_disable_merge_attribute = static_cast<KMemoryBlockDisableMergeAttribute>(
                    m_disable_merge_attribute & ~KMemoryBlockDisableMergeAttribute::IpcLeft);
            }
        }
    }

    constexpr KMemoryBlockDisableMergeAttribute GetDisableMergeAttribute() const {
        return m_disable_merge_attribute;
    }
};
static_assert(std::is_trivially_destructible<KMemoryBlock>::value);

} // namespace Kernel


code:
// SPDX-FileCopyrightText: Copyright 2020 yuzu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#include "core/hle/kernel/k_memory_block_manager.h"

namespace Kernel {

KMemoryBlockManager::KMemoryBlockManager() = default;

Result KMemoryBlockManager::Initialize(KProcessAddress st, KProcessAddress nd,
                                       KMemoryBlockSlabManager* slab_manager) {
    // Allocate a block to encapsulate the address space, insert it into the tree.
    KMemoryBlock* start_block = slab_manager->Allocate();
    R_UNLESS(start_block != nullptr, ResultOutOfResource);

    // Set our start and end.
    m_start_address = st;
    m_end_address = nd;
    ASSERT(Common::IsAligned(GetInteger(m_start_address), PageSize));
    ASSERT(Common::IsAligned(GetInteger(m_end_address), PageSize));

    // Initialize and insert the block.
    start_block->Initialize(m_start_address, (m_end_address - m_start_address) / PageSize,
                            KMemoryState::Free, KMemoryPermission::None, KMemoryAttribute::None);
    m_memory_block_tree.insert(*start_block);

    R_SUCCEED();
}

void KMemoryBlockManager::Finalize(KMemoryBlockSlabManager* slab_manager,
                                   BlockCallback&& block_callback) {
    // Erase every block until we have none left.
    auto it = m_memory_block_tree.begin();
    while (it != m_memory_block_tree.end()) {
        KMemoryBlock* block = std::addressof(*it);
        it = m_memory_block_tree.erase(it);
        block_callback(block->GetAddress(), block->GetSize());
        slab_manager->Free(block);
    }

    ASSERT(m_memory_block_tree.empty());
}

KProcessAddress KMemoryBlockManager::FindFreeArea(KProcessAddress region_start,
                                                  size_t region_num_pages, size_t num_pages,
                                                  size_t alignment, size_t offset,
                                                  size_t guard_pages) const {
    if (num_pages > 0) {
        const KProcessAddress region_end = region_start + region_num_pages * PageSize;
        const KProcessAddress region_last = region_end - 1;
        for (const_iterator it = this->FindIterator(region_start); it != m_memory_block_tree.cend();
             it++) {
            const KMemoryInfo info = it->GetMemoryInfo();
            if (region_last < info.GetAddress()) {
                break;
            }
            if (info.m_state != KMemoryState::Free) {
                continue;
            }

            KProcessAddress area =
                (info.GetAddress() <= GetInteger(region_start)) ? region_start : info.GetAddress();
            area += guard_pages * PageSize;

            const KProcessAddress offset_area =
                Common::AlignDown(GetInteger(area), alignment) + offset;
            area = (area <= offset_area) ? offset_area : offset_area + alignment;

            const KProcessAddress area_end = area + num_pages * PageSize + guard_pages * PageSize;
            const KProcessAddress area_last = area_end - 1;

            if (info.GetAddress() <= GetInteger(area) && area < area_last &&
                area_last <= region_last && area_last <= info.GetLastAddress()) {
                return area;
            }
        }
    }

    return {};
}

void KMemoryBlockManager::CoalesceForUpdate(KMemoryBlockManagerUpdateAllocator* allocator,
                                            KProcessAddress address, size_t num_pages) {
    // Find the iterator now that we've updated.
    iterator it = this->FindIterator(address);
    if (address != m_start_address) {
        it--;
    }

    // Coalesce blocks that we can.
    while (true) {
        iterator prev = it++;
        if (it == m_memory_block_tree.end()) {
            break;
        }

        if (prev->CanMergeWith(*it)) {
            KMemoryBlock* block = std::addressof(*it);
            m_memory_block_tree.erase(it);
            prev->Add(*block);
            allocator->Free(block);
            it = prev;
        }

        if (address + num_pages * PageSize < it->GetMemoryInfo().GetEndAddress()) {
            break;
        }
    }
}

void KMemoryBlockManager::Update(KMemoryBlockManagerUpdateAllocator* allocator,
                                 KProcessAddress address, size_t num_pages, KMemoryState state,
                                 KMemoryPermission perm, KMemoryAttribute attr,
                                 KMemoryBlockDisableMergeAttribute set_disable_attr,
                                 KMemoryBlockDisableMergeAttribute clear_disable_attr) {
    // Ensure for auditing that we never end up with an invalid tree.
    KScopedMemoryBlockManagerAuditor auditor(this);
    ASSERT(Common::IsAligned(GetInteger(address), PageSize));
    ASSERT((attr & (KMemoryAttribute::IpcLocked | KMemoryAttribute::DeviceShared)) ==
           KMemoryAttribute::None);

    KProcessAddress cur_address = address;
    size_t remaining_pages = num_pages;
    iterator it = this->FindIterator(address);

    while (remaining_pages > 0) {
        const size_t remaining_size = remaining_pages * PageSize;
        KMemoryInfo cur_info = it->GetMemoryInfo();
        if (it->HasProperties(state, perm, attr)) {
            // If we already have the right properties, just advance.
            if (cur_address + remaining_size < cur_info.GetEndAddress()) {
                remaining_pages = 0;
                cur_address += remaining_size;
            } else {
                remaining_pages =
                    (cur_address + remaining_size - cur_info.GetEndAddress()) / PageSize;
                cur_address = cur_info.GetEndAddress();
            }
        } else {
            // If we need to, create a new block before and insert it.
            if (cur_info.GetAddress() != cur_address) {
                KMemoryBlock* new_block = allocator->Allocate();

                it->Split(new_block, cur_address);
                it = m_memory_block_tree.insert(*new_block);
                it++;

                cur_info = it->GetMemoryInfo();
                cur_address = cur_info.GetAddress();
            }

            // If we need to, create a new block after and insert it.
            if (cur_info.GetSize() > remaining_size) {
                KMemoryBlock* new_block = allocator->Allocate();

                it->Split(new_block, cur_address + remaining_size);
                it = m_memory_block_tree.insert(*new_block);

                cur_info = it->GetMemoryInfo();
            }

            // Update block state.
            it->Update(state, perm, attr, it->GetAddress() == address,
                       static_cast<u8>(set_disable_attr), static_cast<u8>(clear_disable_attr));
            cur_address += cur_info.GetSize();
            remaining_pages -= cur_info.GetNumPages();
        }
        it++;
    }

    this->CoalesceForUpdate(allocator, address, num_pages);
}

void KMemoryBlockManager::UpdateIfMatch(KMemoryBlockManagerUpdateAllocator* allocator,
                                        KProcessAddress address, size_t num_pages,
                                        KMemoryState test_state, KMemoryPermission test_perm,
                                        KMemoryAttribute test_attr, KMemoryState state,
                                        KMemoryPermission perm, KMemoryAttribute attr,
                                        KMemoryBlockDisableMergeAttribute set_disable_attr,
                                        KMemoryBlockDisableMergeAttribute clear_disable_attr) {
    // Ensure for auditing that we never end up with an invalid tree.
    KScopedMemoryBlockManagerAuditor auditor(this);
    ASSERT(Common::IsAligned(GetInteger(address), PageSize));
    ASSERT((attr & (KMemoryAttribute::IpcLocked | KMemoryAttribute::DeviceShared)) ==
           KMemoryAttribute::None);

    KProcessAddress cur_address = address;
    size_t remaining_pages = num_pages;
    iterator it = this->FindIterator(address);

    while (remaining_pages > 0) {
        const size_t remaining_size = remaining_pages * PageSize;
        KMemoryInfo cur_info = it->GetMemoryInfo();
        if (it->HasProperties(test_state, test_perm, test_attr) &&
            !it->HasProperties(state, perm, attr)) {
            // If we need to, create a new block before and insert it.
            if (cur_info.GetAddress() != cur_address) {
                KMemoryBlock* new_block = allocator->Allocate();

                it->Split(new_block, cur_address);
                it = m_memory_block_tree.insert(*new_block);
                it++;

                cur_info = it->GetMemoryInfo();
                cur_address = cur_info.GetAddress();
            }

            // If we need to, create a new block after and insert it.
            if (cur_info.GetSize() > remaining_size) {
                KMemoryBlock* new_block = allocator->Allocate();

                it->Split(new_block, cur_address + remaining_size);
                it = m_memory_block_tree.insert(*new_block);

                cur_info = it->GetMemoryInfo();
            }

            // Update block state.
            it->Update(state, perm, attr, false, static_cast<u8>(set_disable_attr),
                       static_cast<u8>(clear_disable_attr));
            cur_address += cur_info.GetSize();
            remaining_pages -= cur_info.GetNumPages();
        } else {
            // If we already have the right properties, just advance.
            if (cur_address + remaining_size < cur_info.GetEndAddress()) {
                remaining_pages = 0;
                cur_address += remaining_size;
            } else {
                remaining_pages =
                    (cur_address + remaining_size - cur_info.GetEndAddress()) / PageSize;
                cur_address = cur_info.GetEndAddress();
            }
        }
        it++;
    }

    this->CoalesceForUpdate(allocator, address, num_pages);
}

void KMemoryBlockManager::UpdateLock(KMemoryBlockManagerUpdateAllocator* allocator,
                                     KProcessAddress address, size_t num_pages,
                                     MemoryBlockLockFunction lock_func, KMemoryPermission perm) {
    // Ensure for auditing that we never end up with an invalid tree.
    KScopedMemoryBlockManagerAuditor auditor(this);
    ASSERT(Common::IsAligned(GetInteger(address), PageSize));

    KProcessAddress cur_address = address;
    size_t remaining_pages = num_pages;
    iterator it = this->FindIterator(address);

    const KProcessAddress end_address = address + (num_pages * PageSize);

    while (remaining_pages > 0) {
        const size_t remaining_size = remaining_pages * PageSize;
        KMemoryInfo cur_info = it->GetMemoryInfo();

        // If we need to, create a new block before and insert it.
        if (cur_info.m_address != cur_address) {
            KMemoryBlock* new_block = allocator->Allocate();

            it->Split(new_block, cur_address);
            it = m_memory_block_tree.insert(*new_block);
            it++;

            cur_info = it->GetMemoryInfo();
            cur_address = cur_info.GetAddress();
        }

        if (cur_info.GetSize() > remaining_size) {
            // If we need to, create a new block after and insert it.
            KMemoryBlock* new_block = allocator->Allocate();

            it->Split(new_block, cur_address + remaining_size);
            it = m_memory_block_tree.insert(*new_block);

            cur_info = it->GetMemoryInfo();
        }

        // Call the locked update function.
        (std::addressof(*it)->*lock_func)(perm, cur_info.GetAddress() == address,
                                          cur_info.GetEndAddress() == end_address);
        cur_address += cur_info.GetSize();
        remaining_pages -= cur_info.GetNumPages();
        it++;
    }

    this->CoalesceForUpdate(allocator, address, num_pages);
}

void KMemoryBlockManager::UpdateAttribute(KMemoryBlockManagerUpdateAllocator* allocator,
                                          KProcessAddress address, size_t num_pages,
                                          KMemoryAttribute mask, KMemoryAttribute attr) {
    // Ensure for auditing that we never end up with an invalid tree.
    KScopedMemoryBlockManagerAuditor auditor(this);
    ASSERT(Common::IsAligned(GetInteger(address), PageSize));

    KProcessAddress cur_address = address;
    size_t remaining_pages = num_pages;
    iterator it = this->FindIterator(address);

    while (remaining_pages > 0) {
        const size_t remaining_size = remaining_pages * PageSize;
        KMemoryInfo cur_info = it->GetMemoryInfo();

        if ((it->GetAttribute() & mask) != attr) {
            // If we need to, create a new block before and insert it.
            if (cur_info.GetAddress() != GetInteger(cur_address)) {
                KMemoryBlock* new_block = allocator->Allocate();

                it->Split(new_block, cur_address);
                it = m_memory_block_tree.insert(*new_block);
                it++;

                cur_info = it->GetMemoryInfo();
                cur_address = cur_info.GetAddress();
            }

            // If we need to, create a new block after and insert it.
            if (cur_info.GetSize() > remaining_size) {
                KMemoryBlock* new_block = allocator->Allocate();

                it->Split(new_block, cur_address + remaining_size);
                it = m_memory_block_tree.insert(*new_block);

                cur_info = it->GetMemoryInfo();
            }

            // Update block state.
            it->UpdateAttribute(mask, attr);
            cur_address += cur_info.GetSize();
            remaining_pages -= cur_info.GetNumPages();
        } else {
            // If we already have the right attributes, just advance.
            if (cur_address + remaining_size < cur_info.GetEndAddress()) {
                remaining_pages = 0;
                cur_address += remaining_size;
            } else {
                remaining_pages =
                    (cur_address + remaining_size - cur_info.GetEndAddress()) / PageSize;
                cur_address = cur_info.GetEndAddress();
            }
        }
        it++;
    }

    this->CoalesceForUpdate(allocator, address, num_pages);
}

// Debug.
bool KMemoryBlockManager::CheckState() const {
    // Loop over every block, ensuring that we are sorted and coalesced.
    auto it = m_memory_block_tree.cbegin();
    auto prev = it++;
    while (it != m_memory_block_tree.cend()) {
        const KMemoryInfo prev_info = prev->GetMemoryInfo();
        const KMemoryInfo cur_info = it->GetMemoryInfo();

        // Sequential blocks which can be merged should be merged.
        if (prev->CanMergeWith(*it)) {
            return false;
        }

        // Sequential blocks should be sequential.
        if (prev_info.GetEndAddress() != cur_info.GetAddress()) {
            return false;
        }

        // If the block is ipc locked, it must have a count.
        if ((cur_info.m_attribute & KMemoryAttribute::IpcLocked) != KMemoryAttribute::None &&
            cur_info.m_ipc_lock_count == 0) {
            return false;
        }

        // If the block is device shared, it must have a count.
        if ((cur_info.m_attribute & KMemoryAttribute::DeviceShared) != KMemoryAttribute::None &&
            cur_info.m_device_use_count == 0) {
            return false;
        }

        // Advance the iterator.
        prev = it++;
    }

    // Our loop will miss checking the last block, potentially, so check it.
    if (prev != m_memory_block_tree.cend()) {
        const KMemoryInfo prev_info = prev->GetMemoryInfo();
        // If the block is ipc locked, it must have a count.
        if ((prev_info.m_attribute & KMemoryAttribute::IpcLocked) != KMemoryAttribute::None &&
            prev_info.m_ipc_lock_count == 0) {
            return false;
        }

        // If the block is device shared, it must have a count.
        if ((prev_info.m_attribute & KMemoryAttribute::DeviceShared) != KMemoryAttribute::None &&
            prev_info.m_device_use_count == 0) {
            return false;
        }
    }

    return true;
}

} // namespace Kernel

repo_name: litucks/torzu

license: mit


code:
// SPDX-FileCopyrightText: Copyright 2020 yuzu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#pragma once

#include <array>
#include <functional>

#include "common/common_funcs.h"
#include "core/hle/kernel/k_dynamic_resource_manager.h"
#include "core/hle/kernel/k_memory_block.h"
#include "core/hle/kernel/k_typed_address.h"

namespace Kernel {

class KMemoryBlockManagerUpdateAllocator {
public:
    static constexpr size_t MaxBlocks = 2;

private:
    std::array<KMemoryBlock*, MaxBlocks> m_blocks{};
    size_t m_index{MaxBlocks};
    KMemoryBlockSlabManager* m_slab_manager{};

private:
    Result Initialize(size_t num_blocks) {
        // Check num blocks.
        ASSERT(num_blocks <= MaxBlocks);

        // Set index.
        m_index = MaxBlocks - num_blocks;

        // Allocate the blocks.
        for (size_t i = 0; i < num_blocks && i < MaxBlocks; ++i) {
            m_blocks[m_index + i] = m_slab_manager->Allocate();
            R_UNLESS(m_blocks[m_index + i] != nullptr, ResultOutOfResource);
        }

        R_SUCCEED();
    }

public:
    KMemoryBlockManagerUpdateAllocator(Result* out_result, KMemoryBlockSlabManager* sm,
                                       size_t num_blocks = MaxBlocks)
        : m_slab_manager(sm) {
        *out_result = this->Initialize(num_blocks);
    }

    ~KMemoryBlockManagerUpdateAllocator() {
        for (const auto& block : m_blocks) {
            if (block != nullptr) {
                m_slab_manager->Free(block);
            }
        }
    }

    KMemoryBlock* Allocate() {
        ASSERT(m_index < MaxBlocks);
        ASSERT(m_blocks[m_index] != nullptr);
        KMemoryBlock* block = nullptr;
        std::swap(block, m_blocks[m_index++]);
        return block;
    }

    void Free(KMemoryBlock* block) {
        ASSERT(m_index <= MaxBlocks);
        ASSERT(block != nullptr);
        if (m_index == 0) {
            m_slab_manager->Free(block);
        } else {
            m_blocks[--m_index] = block;
        }
    }
};

class KMemoryBlockManager final {
public:
    using MemoryBlockTree =
        Common::IntrusiveRedBlackTreeBaseTraits<KMemoryBlock>::TreeType<KMemoryBlock>;
    using MemoryBlockLockFunction = void (KMemoryBlock::*)(KMemoryPermission new_perm, bool left,
                                                           bool right);
    using iterator = MemoryBlockTree::iterator;
    using const_iterator = MemoryBlockTree::const_iterator;

public:
    KMemoryBlockManager();

    using BlockCallback = std::function<void(Common::ProcessAddress, u64)>;

    Result Initialize(KProcessAddress st, KProcessAddress nd,
                      KMemoryBlockSlabManager* slab_manager);
    void Finalize(KMemoryBlockSlabManager* slab_manager, BlockCallback&& block_callback);

    iterator end() {
        return m_memory_block_tree.end();
    }
    const_iterator end() const {
        return m_memory_block_tree.end();
    }
    const_iterator cend() const {
        return m_memory_block_tree.cend();
    }

    KProcessAddress FindFreeArea(KProcessAddress region_start, size_t region_num_pages,
                                 size_t num_pages, size_t alignment, size_t offset,
                                 size_t guard_pages) const;

    void Update(KMemoryBlockManagerUpdateAllocator* allocator, KProcessAddress address,
                size_t num_pages, KMemoryState state, KMemoryPermission perm, KMemoryAttribute attr,
                KMemoryBlockDisableMergeAttribute set_disable_attr,
                KMemoryBlockDisableMergeAttribute clear_disable_attr);
    void UpdateLock(KMemoryBlockManagerUpdateAllocator* allocator, KProcessAddress address,
                    size_t num_pages, MemoryBlockLockFunction lock_func, KMemoryPermission perm);

    void UpdateIfMatch(KMemoryBlockManagerUpdateAllocator* allocator, KProcessAddress address,
                       size_t num_pages, KMemoryState test_state, KMemoryPermission test_perm,
                       KMemoryAttribute test_attr, KMemoryState state, KMemoryPermission perm,
                       KMemoryAttribute attr, KMemoryBlockDisableMergeAttribute set_disable_attr,
                       KMemoryBlockDisableMergeAttribute clear_disable_attr);

    void UpdateAttribute(KMemoryBlockManagerUpdateAllocator* allocator, KProcessAddress address,
                         size_t num_pages, KMemoryAttribute mask, KMemoryAttribute attr);

    iterator FindIterator(KProcessAddress address) const {
        return m_memory_block_tree.find(KMemoryBlock(
            address, 1, KMemoryState::Free, KMemoryPermission::None, KMemoryAttribute::None));
    }

    const KMemoryBlock* FindBlock(KProcessAddress address) const {
        if (const_iterator it = this->FindIterator(address); it != m_memory_block_tree.end()) {
            return std::addressof(*it);
        }

        return nullptr;
    }

    // Debug.
    bool CheckState() const;

private:
    void CoalesceForUpdate(KMemoryBlockManagerUpdateAllocator* allocator, KProcessAddress address,
                           size_t num_pages);

    MemoryBlockTree m_memory_block_tree;
    KProcessAddress m_start_address{};
    KProcessAddress m_end_address{};
};

class KScopedMemoryBlockManagerAuditor {
public:
    explicit KScopedMemoryBlockManagerAuditor(KMemoryBlockManager* m) : m_manager(m) {}
    explicit KScopedMemoryBlockManagerAuditor(KMemoryBlockManager& m)
        : KScopedMemoryBlockManagerAuditor(std::addressof(m)) {}
    ~KScopedMemoryBlockManagerAuditor() = default;

private:
    KMemoryBlockManager* m_manager;
};

} // namespace Kernel


code:
// SPDX-FileCopyrightText: Copyright 2021 yuzu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#include <array>

#include "common/alignment.h"
#include "core/hle/kernel/k_memory_layout.h"
#include "core/hle/kernel/k_system_control.h"

namespace Kernel {

namespace {

template <typename... Args>
KMemoryRegion* AllocateRegion(KMemoryRegionAllocator& memory_region_allocator, Args&&... args) {
    return memory_region_allocator.Allocate(std::forward<Args>(args)...);
}

} // namespace

KMemoryRegionTree::KMemoryRegionTree(KMemoryRegionAllocator& memory_region_allocator)
    : m_memory_region_allocator{memory_region_allocator} {}

void KMemoryRegionTree::InsertDirectly(u64 address, u64 last_address, u32 attr, u32 type_id) {
    this->insert(*AllocateRegion(m_memory_region_allocator, address, last_address, attr, type_id));
}

bool KMemoryRegionTree::Insert(u64 address, size_t size, u32 type_id, u32 new_attr, u32 old_attr) {
    // Locate the memory region that contains the address.
    KMemoryRegion* found = this->FindModifiable(address);

    // We require that the old attr is correct.
    if (found->GetAttributes() != old_attr) {
        return false;
    }

    // We further require that the region can be split from the old region.
    const u64 inserted_region_end = address + size;
    const u64 inserted_region_last = inserted_region_end - 1;
    if (found->GetLastAddress() < inserted_region_last) {
        return false;
    }

    // Further, we require that the type id is a valid transformation.
    if (!found->CanDerive(type_id)) {
        return false;
    }

    // Cache information from the region before we remove it.
    const u64 old_address = found->GetAddress();
    const u64 old_last = found->GetLastAddress();
    const u64 old_pair = found->GetPairAddress();
    const u32 old_type = found->GetType();

    // Erase the existing region from the tree.
    this->erase(this->iterator_to(*found));

    // Insert the new region into the tree.
    if (old_address == address) {
        // Reuse the old object for the new region, if we can.
        found->Reset(address, inserted_region_last, old_pair, new_attr, type_id);
        this->insert(*found);
    } else {
        // If we can't reuse, adjust the old region.
        found->Reset(old_address, address - 1, old_pair, old_attr, old_type);
        this->insert(*found);

        // Insert a new region for the split.
        const u64 new_pair = (old_pair != std::numeric_limits<u64>::max())
                                 ? old_pair + (address - old_address)
                                 : old_pair;
        this->insert(*AllocateRegion(m_memory_region_allocator, address, inserted_region_last,
                                     new_pair, new_attr, type_id));
    }

    // If we need to insert a region after the region, do so.
    if (old_last != inserted_region_last) {
        const u64 after_pair = (old_pair != std::numeric_limits<u64>::max())
                                   ? old_pair + (inserted_region_end - old_address)
                                   : old_pair;
        this->insert(*AllocateRegion(m_memory_region_allocator, inserted_region_end, old_last,
                                     after_pair, old_attr, old_type));
    }

    return true;
}

KVirtualAddress KMemoryRegionTree::GetRandomAlignedRegion(size_t size, size_t alignment,
                                                          u32 type_id) {
    // We want to find the total extents of the type id.
    const auto extents = this->GetDerivedRegionExtents(static_cast<KMemoryRegionType>(type_id));

    // Ensure that our alignment is correct.
    ASSERT(Common::IsAligned(extents.GetAddress(), alignment));

    const u64 first_address = extents.GetAddress();
    const u64 last_address = extents.GetLastAddress();

    const u64 first_index = first_address / alignment;
    const u64 last_index = last_address / alignment;

    while (true) {
        const u64 candidate =
            KSystemControl::GenerateRandomRange(first_index, last_index) * alignment;

        // Ensure that the candidate doesn't overflow with the size.
        if (!(candidate < candidate + size)) {
            continue;
        }

        const u64 candidate_last = candidate + size - 1;

        // Ensure that the candidate fits within the region.
        if (candidate_last > last_address) {
            continue;
        }

        // Locate the candidate region, and ensure it fits and has the correct type id.
        if (const auto& candidate_region = *this->Find(candidate);
            !(candidate_last <= candidate_region.GetLastAddress() &&
              candidate_region.GetType() == type_id)) {
            continue;
        }

        return candidate;
    }
}

KMemoryLayout::KMemoryLayout()
    : m_virtual_tree{m_memory_region_allocator}, m_physical_tree{m_memory_region_allocator},
      m_virtual_linear_tree{m_memory_region_allocator}, m_physical_linear_tree{
                                                            m_memory_region_allocator} {}

void KMemoryLayout::InitializeLinearMemoryRegionTrees(KPhysicalAddress aligned_linear_phys_start,
                                                      KVirtualAddress linear_virtual_start) {
    // Set static differences.
    m_linear_phys_to_virt_diff =
        GetInteger(linear_virtual_start) - GetInteger(aligned_linear_phys_start);
    m_linear_virt_to_phys_diff =
        GetInteger(aligned_linear_phys_start) - GetInteger(linear_virtual_start);

    // Initialize linear trees.
    for (auto& region : GetPhysicalMemoryRegionTree()) {
        if (region.HasTypeAttribute(KMemoryRegionAttr_LinearMapped)) {
            GetPhysicalLinearMemoryRegionTree().InsertDirectly(
                region.GetAddress(), region.GetLastAddress(), region.GetAttributes(),
                region.GetType());
        }
    }

    for (auto& region : GetVirtualMemoryRegionTree()) {
        if (region.IsDerivedFrom(KMemoryRegionType_Dram)) {
            GetVirtualLinearMemoryRegionTree().InsertDirectly(
                region.GetAddress(), region.GetLastAddress(), region.GetAttributes(),
                region.GetType());
        }
    }
}

size_t KMemoryLayout::GetResourceRegionSizeForInit(bool use_extra_resource) {
    return KernelResourceSize + KSystemControl::SecureAppletMemorySize +
           (use_extra_resource ? KernelSlabHeapAdditionalSize + KernelPageBufferAdditionalSize : 0);
}

} // namespace Kernel


code:
// SPDX-FileCopyrightText: Copyright 2021 yuzu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#pragma once

#include <utility>

#include "common/alignment.h"
#include "common/literals.h"
#include "core/device_memory.h"
#include "core/hle/kernel/k_memory_region.h"
#include "core/hle/kernel/k_memory_region_type.h"
#include "core/hle/kernel/k_typed_address.h"
#include "core/hle/kernel/memory_types.h"

namespace Kernel {

using namespace Common::Literals;

constexpr std::size_t L1BlockSize = 1_GiB;
constexpr std::size_t L2BlockSize = 2_MiB;

constexpr std::size_t GetMaximumOverheadSize(std::size_t size) {
    return (Common::DivideUp(size, L1BlockSize) + Common::DivideUp(size, L2BlockSize)) * PageSize;
}

constexpr std::size_t MainMemorySize = 4_GiB;
constexpr std::size_t MainMemorySizeMax = 8_GiB;

constexpr std::size_t ReservedEarlyDramSize = 384_KiB;
constexpr std::size_t DramPhysicalAddress = 0x80000000;

constexpr std::size_t KernelAslrAlignment = 2_MiB;
constexpr std::size_t KernelVirtualAddressSpaceWidth = 1ULL << 39;
constexpr std::size_t KernelPhysicalAddressSpaceWidth = 1ULL << 48;

constexpr std::size_t KernelVirtualAddressSpaceBase = 0ULL - KernelVirtualAddressSpaceWidth;
constexpr std::size_t KernelVirtualAddressSpaceEnd =
    KernelVirtualAddressSpaceBase + (KernelVirtualAddressSpaceWidth - KernelAslrAlignment);
constexpr std::size_t KernelVirtualAddressSpaceLast = KernelVirtualAddressSpaceEnd - 1ULL;
constexpr std::size_t KernelVirtualAddressSpaceSize =
    KernelVirtualAddressSpaceEnd - KernelVirtualAddressSpaceBase;
constexpr std::size_t KernelVirtualAddressCodeBase = KernelVirtualAddressSpaceBase;
constexpr std::size_t KernelVirtualAddressCodeSize = 392_KiB;
constexpr std::size_t KernelVirtualAddressCodeEnd =
    KernelVirtualAddressCodeBase + KernelVirtualAddressCodeSize;

constexpr std::size_t KernelPhysicalAddressSpaceBase = 0ULL;
constexpr std::size_t KernelPhysicalAddressSpaceEnd =
    KernelPhysicalAddressSpaceBase + KernelPhysicalAddressSpaceWidth;
constexpr std::size_t KernelPhysicalAddressSpaceLast = KernelPhysicalAddressSpaceEnd - 1ULL;
constexpr std::size_t KernelPhysicalAddressSpaceSize =
    KernelPhysicalAddressSpaceEnd - KernelPhysicalAddressSpaceBase;
constexpr std::size_t KernelPhysicalAddressCodeBase = DramPhysicalAddress + ReservedEarlyDramSize;

constexpr std::size_t KernelPageTableHeapSize = GetMaximumOverheadSize(MainMemorySizeMax);
constexpr std::size_t KernelInitialPageHeapSize = 128_KiB;

constexpr std::size_t KernelSlabHeapDataSize = 5_MiB;
constexpr std::size_t KernelSlabHeapGapsSizeMax = 2_MiB - 64_KiB;
constexpr std::size_t KernelSlabHeapSize = KernelSlabHeapDataSize + KernelSlabHeapGapsSizeMax;

// NOTE: This is calculated from KThread slab counts, assuming KThread size <= 0x860.
constexpr size_t KernelPageBufferHeapSize = 0x3E0000;
constexpr size_t KernelSlabHeapAdditionalSize = 0x148000;
constexpr size_t KernelPageBufferAdditionalSize = 0x33C000;

constexpr std::size_t KernelResourceSize = KernelPageTableHeapSize + KernelInitialPageHeapSize +
                                           KernelSlabHeapSize + KernelPageBufferHeapSize;

//! NB: Use KThread::GetAddressKeyIsKernel().
//! See explanation for deviation of GetAddressKey.
bool IsKernelAddressKey(KProcessAddress key) = delete;

constexpr bool IsKernelAddress(KProcessAddress address) {
    return KernelVirtualAddressSpaceBase <= GetInteger(address) &&
           address < KernelVirtualAddressSpaceEnd;
}

class KMemoryLayout final {
public:
    KMemoryLayout();

    KMemoryRegionTree& GetVirtualMemoryRegionTree() {
        return m_virtual_tree;
    }
    const KMemoryRegionTree& GetVirtualMemoryRegionTree() const {
        return m_virtual_tree;
    }
    KMemoryRegionTree& GetPhysicalMemoryRegionTree() {
        return m_physical_tree;
    }
    const KMemoryRegionTree& GetPhysicalMemoryRegionTree() const {
        return m_physical_tree;
    }
    KMemoryRegionTree& GetVirtualLinearMemoryRegionTree() {
        return m_virtual_linear_tree;
    }
    const KMemoryRegionTree& GetVirtualLinearMemoryRegionTree() const {
        return m_virtual_linear_tree;
    }
    KMemoryRegionTree& GetPhysicalLinearMemoryRegionTree() {
        return m_physical_linear_tree;
    }
    const KMemoryRegionTree& GetPhysicalLinearMemoryRegionTree() const {
        return m_physical_linear_tree;
    }

    KVirtualAddress GetLinearVirtualAddress(KPhysicalAddress address) const {
        return GetInteger(address) + m_linear_phys_to_virt_diff;
    }
    KPhysicalAddress GetLinearPhysicalAddress(KVirtualAddress address) const {
        return GetInteger(address) + m_linear_virt_to_phys_diff;
    }

    const KMemoryRegion* FindVirtual(KVirtualAddress address) const {
        return Find(address, GetVirtualMemoryRegionTree());
    }
    const KMemoryRegion* FindPhysical(KPhysicalAddress address) const {
        return Find(address, GetPhysicalMemoryRegionTree());
    }

    const KMemoryRegion* FindVirtualLinear(KVirtualAddress address) const {
        return Find(address, GetVirtualLinearMemoryRegionTree());
    }
    const KMemoryRegion* FindPhysicalLinear(KPhysicalAddress address) const {
        return Find(address, GetPhysicalLinearMemoryRegionTree());
    }

    KVirtualAddress GetMainStackTopAddress(s32 core_id) const {
        return GetStackTopAddress(core_id, KMemoryRegionType_KernelMiscMainStack);
    }
    KVirtualAddress GetIdleStackTopAddress(s32 core_id) const {
        return GetStackTopAddress(core_id, KMemoryRegionType_KernelMiscIdleStack);
    }
    KVirtualAddress GetExceptionStackTopAddress(s32 core_id) const {
        return GetStackTopAddress(core_id, KMemoryRegionType_KernelMiscExceptionStack);
    }

    const KMemoryRegion& GetSlabRegion() const {
        return Dereference(GetVirtualMemoryRegionTree().FindByType(KMemoryRegionType_KernelSlab));
    }
    const KMemoryRegion& GetDeviceRegion(KMemoryRegionType type) const {
        return Dereference(GetPhysicalMemoryRegionTree().FindFirstDerived(type));
    }
    KPhysicalAddress GetDevicePhysicalAddress(KMemoryRegionType type) const {
        return GetDeviceRegion(type).GetAddress();
    }
    KVirtualAddress GetDeviceVirtualAddress(KMemoryRegionType type) const {
        return GetDeviceRegion(type).GetPairAddress();
    }

    const KMemoryRegion& GetPoolManagementRegion() const {
        return Dereference(
            GetVirtualMemoryRegionTree().FindByType(KMemoryRegionType_VirtualDramPoolManagement));
    }
    const KMemoryRegion& GetPageTableHeapRegion() const {
        return Dereference(
            GetVirtualMemoryRegionTree().FindByType(KMemoryRegionType_VirtualDramKernelPtHeap));
    }
    const KMemoryRegion& GetKernelStackRegion() const {
        return Dereference(GetVirtualMemoryRegionTree().FindByType(KMemoryRegionType_KernelStack));
    }
    const KMemoryRegion& GetTempRegion() const {
        return Dereference(GetVirtualMemoryRegionTree().FindByType(KMemoryRegionType_KernelTemp));
    }

    const KMemoryRegion& GetKernelTraceBufferRegion() const {
        return Dereference(GetVirtualLinearMemoryRegionTree().FindByType(
            KMemoryRegionType_VirtualDramKernelTraceBuffer));
    }

    const KMemoryRegion& GetSecureAppletMemoryRegion() {
        return Dereference(GetVirtualMemoryRegionTree().FindByType(
            KMemoryRegionType_VirtualDramKernelSecureAppletMemory));
    }

    const KMemoryRegion& GetVirtualLinearRegion(KVirtualAddress address) const {
        return Dereference(FindVirtualLinear(address));
    }

    const KMemoryRegion& GetPhysicalLinearRegion(KPhysicalAddress address) const {
        return Dereference(FindPhysicalLinear(address));
    }

    const KMemoryRegion* GetPhysicalKernelTraceBufferRegion() const {
        return GetPhysicalMemoryRegionTree().FindFirstDerived(KMemoryRegionType_KernelTraceBuffer);
    }
    const KMemoryRegion* GetPhysicalOnMemoryBootImageRegion() const {
        return GetPhysicalMemoryRegionTree().FindFirstDerived(KMemoryRegionType_OnMemoryBootImage);
    }
    const KMemoryRegion* GetPhysicalDTBRegion() const {
        return GetPhysicalMemoryRegionTree().FindFirstDerived(KMemoryRegionType_DTB);
    }

    bool IsHeapPhysicalAddress(const KMemoryRegion*& region, KPhysicalAddress address) const {
        return IsTypedAddress(region, address, GetPhysicalLinearMemoryRegionTree(),
                              KMemoryRegionType_DramUserPool);
    }
    bool IsHeapVirtualAddress(const KMemoryRegion*& region, KVirtualAddress address) const {
        return IsTypedAddress(region, address, GetVirtualLinearMemoryRegionTree(),
                              KMemoryRegionType_VirtualDramUserPool);
    }

    bool IsHeapPhysicalAddress(const KMemoryRegion*& region, KPhysicalAddress address,
                               size_t size) const {
        return IsTypedAddress(region, address, size, GetPhysicalLinearMemoryRegionTree(),
                              KMemoryRegionType_DramUserPool);
    }
    bool IsHeapVirtualAddress(const KMemoryRegion*& region, KVirtualAddress address,
                              size_t size) const {
        return IsTypedAddress(region, address, size, GetVirtualLinearMemoryRegionTree(),
                              KMemoryRegionType_VirtualDramUserPool);
    }

    bool IsLinearMappedPhysicalAddress(const KMemoryRegion*& region,
                                       KPhysicalAddress address) const {
        return IsTypedAddress(region, address, GetPhysicalLinearMemoryRegionTree(),
                              static_cast<KMemoryRegionType>(KMemoryRegionAttr_LinearMapped));
    }
    bool IsLinearMappedPhysicalAddress(const KMemoryRegion*& region, KPhysicalAddress address,
                                       size_t size) const {
        return IsTypedAddress(region, address, size, GetPhysicalLinearMemoryRegionTree(),
                              static_cast<KMemoryRegionType>(KMemoryRegionAttr_LinearMapped));
    }

    std::pair<size_t, size_t> GetTotalAndKernelMemorySizes() const {
        size_t total_size = 0, kernel_size = 0;
        for (const auto& region : GetPhysicalMemoryRegionTree()) {
            if (region.IsDerivedFrom(KMemoryRegionType_Dram)) {
                total_size += region.GetSize();
                if (!region.IsDerivedFrom(KMemoryRegionType_DramUserPool)) {
                    kernel_size += region.GetSize();
                }
            }
        }
        return std::make_pair(total_size, kernel_size);
    }

    void InitializeLinearMemoryRegionTrees(KPhysicalAddress aligned_linear_phys_start,
                                           KVirtualAddress linear_virtual_start);
    static size_t GetResourceRegionSizeForInit(bool use_extra_resource);

    auto GetKernelRegionExtents() const {
        return GetVirtualMemoryRegionTree().GetDerivedRegionExtents(KMemoryRegionType_Kernel);
    }
    auto GetKernelCodeRegionExtents() const {
        return GetVirtualMemoryRegionTree().GetDerivedRegionExtents(KMemoryRegionType_KernelCode);
    }
    auto GetKernelStackRegionExtents() const {
        return GetVirtualMemoryRegionTree().GetDerivedRegionExtents(KMemoryRegionType_KernelStack);
    }
    auto GetKernelMiscRegionExtents() const {
        return GetVirtualMemoryRegionTree().GetDerivedRegionExtents(KMemoryRegionType_KernelMisc);
    }
    auto GetKernelSlabRegionExtents() const {
        return GetVirtualMemoryRegionTree().GetDerivedRegionExtents(KMemoryRegionType_KernelSlab);
    }

    auto GetLinearRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionAttr_LinearMapped);
    }

    auto GetLinearRegionVirtualExtents() const {
        const auto physical = GetLinearRegionPhysicalExtents();
        return KMemoryRegion(GetInteger(GetLinearVirtualAddress(physical.GetAddress())),
                             GetInteger(GetLinearVirtualAddress(physical.GetLastAddress())), 0,
                             KMemoryRegionType_None);
    }

    auto GetMainMemoryPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(KMemoryRegionType_Dram);
    }
    auto GetCarveoutRegionExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionAttr_CarveoutProtected);
    }

    auto GetKernelRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramKernelBase);
    }
    auto GetKernelCodeRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramKernelCode);
    }
    auto GetKernelSlabRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramKernelSlab);
    }
    auto GetKernelSecureAppletMemoryRegionPhysicalExtents() {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramKernelSecureAppletMemory);
    }
    auto GetKernelPageTableHeapRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramKernelPtHeap);
    }
    auto GetKernelInitPageTableRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramKernelInitPt);
    }

    auto GetKernelPoolManagementRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramPoolManagement);
    }
    auto GetKernelPoolPartitionRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramPoolPartition);
    }
    auto GetKernelSystemPoolRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramSystemPool);
    }
    auto GetKernelSystemNonSecurePoolRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramSystemNonSecurePool);
    }
    auto GetKernelAppletPoolRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramAppletPool);
    }
    auto GetKernelApplicationPoolRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_DramApplicationPool);
    }

    auto GetKernelTraceBufferRegionPhysicalExtents() const {
        return GetPhysicalMemoryRegionTree().GetDerivedRegionExtents(
            KMemoryRegionType_KernelTraceBuffer);
    }

private:
    template <typename AddressType>
    static bool IsTypedAddress(const KMemoryRegion*& region, AddressType address,
                               const KMemoryRegionTree& tree, KMemoryRegionType type) {
        // Check if the cached region already contains the address.
        if (region != nullptr && region->Contains(GetInteger(address))) {
            return true;
        }

        // Find the containing region, and update the cache.
        if (const KMemoryRegion* found = tree.Find(GetInteger(address));
            found != nullptr && found->IsDerivedFrom(type)) {
            region = found;
            return true;
        } else {
            return false;
        }
    }

    template <typename AddressType>
    static bool IsTypedAddress(const KMemoryRegion*& region, AddressType address, size_t size,
                               const KMemoryRegionTree& tree, KMemoryRegionType type) {
        // Get the end of the checked region.
        const u64 last_address = GetInteger(address) + size - 1;

        // Walk the tree to verify the region is correct.
        const KMemoryRegion* cur = (region != nullptr && region->Contains(GetInteger(address)))
                                       ? region
                                       : tree.Find(GetInteger(address));
        while (cur != nullptr && cur->IsDerivedFrom(type)) {
            if (last_address <= cur->GetLastAddress()) {
                region = cur;
                return true;
            }

            cur = cur->GetNext();
        }
        return false;
    }

    template <typename AddressType>
    static const KMemoryRegion* Find(AddressType address, const KMemoryRegionTree& tree) {
        return tree.Find(GetInteger(address));
    }

    static KMemoryRegion& Dereference(KMemoryRegion* region) {
        ASSERT(region != nullptr);
        return *region;
    }

    static const KMemoryRegion& Dereference(const KMemoryRegion* region) {
        ASSERT(region != nullptr);
        return *region;
    }

    KVirtualAddress GetStackTopAddress(s32 core_id, KMemoryRegionType type) const {
        const auto& region = Dereference(
            GetVirtualMemoryRegionTree().FindByTypeAndAttribute(type, static_cast<u32>(core_id)));
        ASSERT(region.GetEndAddress() != 0);
        return region.GetEndAddress();
    }

public:
    static const KMemoryRegion* Find(const KMemoryLayout& layout, KVirtualAddress address) {
        return Find(address, layout.GetVirtualMemoryRegionTree());
    }
    static const KMemoryRegion* Find(const KMemoryLayout& layout, KPhysicalAddress address) {
        return Find(address, layout.GetPhysicalMemoryRegionTree());
    }

private:
    u64 m_linear_phys_to_virt_diff{};
    u64 m_linear_virt_to_phys_diff{};
    KMemoryRegionAllocator m_memory_region_allocator;
    KMemoryRegionTree m_virtual_tree;
    KMemoryRegionTree m_physical_tree;
    KMemoryRegionTree m_virtual_linear_tree;
    KMemoryRegionTree m_physical_linear_tree;
};

namespace Init {

// These should be generic, regardless of board.
void SetupPoolPartitionMemoryRegions(KMemoryLayout& memory_layout);

// These may be implemented in a board-specific manner.
void SetupDevicePhysicalMemoryRegions(KMemoryLayout& memory_layout);
void SetupDramPhysicalMemoryRegions(KMemoryLayout& memory_layout);

} // namespace Init

} // namespace Kernel


code:
// SPDX-FileCopyrightText: Copyright 2020 yuzu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#include <algorithm>

#include "common/alignment.h"
#include "common/assert.h"
#include "common/scope_exit.h"
#include "core/core.h"
#include "core/device_memory.h"
#include "core/hle/kernel/initial_process.h"
#include "core/hle/kernel/k_memory_manager.h"
#include "core/hle/kernel/k_page_group.h"
#include "core/hle/kernel/k_page_table.h"
#include "core/hle/kernel/kernel.h"
#include "core/hle/kernel/svc_results.h"

namespace Kernel {

namespace {

constexpr KMemoryManager::Pool GetPoolFromMemoryRegionType(u32 type) {
    if ((type | KMemoryRegionType_DramApplicationPool) == type) {
        return KMemoryManager::Pool::Application;
    } else if ((type | KMemoryRegionType_DramAppletPool) == type) {
        return KMemoryManager::Pool::Applet;
    } else if ((type | KMemoryRegionType_DramSystemPool) == type) {
        return KMemoryManager::Pool::System;
    } else if ((type | KMemoryRegionType_DramSystemNonSecurePool) == type) {
        return KMemoryManager::Pool::SystemNonSecure;
    } else {
        UNREACHABLE_MSG("InvalidMemoryRegionType for conversion to Pool");
    }
}

} // namespace

KMemoryManager::KMemoryManager(Core::System& system)
    : m_system{system}, m_memory_layout{system.Kernel().MemoryLayout()},
      m_pool_locks{
          KLightLock{system.Kernel()},
          KLightLock{system.Kernel()},
          KLightLock{system.Kernel()},
          KLightLock{system.Kernel()},
      } {}

void KMemoryManager::Initialize(KVirtualAddress management_region, size_t management_region_size) {

    // Clear the management region to zero.
    const KVirtualAddress management_region_end = management_region + management_region_size;
    // std::memset(GetVoidPointer(management_region), 0, management_region_size);

    // Reset our manager count.
    m_num_managers = 0;

    // Traverse the virtual memory layout tree, initializing each manager as appropriate.
    while (m_num_managers != MaxManagerCount) {
        // Locate the region that should initialize the current manager.
        KPhysicalAddress region_address = 0;
        size_t region_size = 0;
        Pool region_pool = Pool::Count;
        for (const auto& it : m_system.Kernel().MemoryLayout().GetPhysicalMemoryRegionTree()) {
            // We only care about regions that we need to create managers for.
            if (!it.IsDerivedFrom(KMemoryRegionType_DramUserPool)) {
                continue;
            }

            // We want to initialize the managers in order.
            if (it.GetAttributes() != m_num_managers) {
                continue;
            }

            const KPhysicalAddress cur_start = it.GetAddress();
            const KPhysicalAddress cur_end = it.GetEndAddress();

            // Validate the region.
            ASSERT(cur_end != 0);
            ASSERT(cur_start != 0);
            ASSERT(it.GetSize() > 0);

            // Update the region's extents.
            if (region_address == 0) {
                region_address = cur_start;
                region_size = it.GetSize();
                region_pool = GetPoolFromMemoryRegionType(it.GetType());
            } else {
                ASSERT(cur_start == region_address + region_size);

                // Update the size.
                region_size = cur_end - region_address;
                ASSERT(GetPoolFromMemoryRegionType(it.GetType()) == region_pool);
            }
        }

        // If we didn't find a region, we're done.
        if (region_size == 0) {
            break;
        }

        // Initialize a new manager for the region.
        Impl* manager = std::addressof(m_managers[m_num_managers++]);
        ASSERT(m_num_managers <= m_managers.size());

        const size_t cur_size = manager->Initialize(region_address, region_size, management_region,
                                                    management_region_end, region_pool);
        management_region += cur_size;
        ASSERT(management_region <= management_region_end);

        // Insert the manager into the pool list.
        const auto region_pool_index = static_cast<u32>(region_pool);
        if (m_pool_managers_tail[region_pool_index] == nullptr) {
            m_pool_managers_head[region_pool_index] = manager;
        } else {
            m_pool_managers_tail[region_pool_index]->SetNext(manager);
            manager->SetPrev(m_pool_managers_tail[region_pool_index]);
        }
        m_pool_managers_tail[region_pool_index] = manager;
    }

    // Free each region to its corresponding heap.
    size_t reserved_sizes[MaxManagerCount] = {};
    const KPhysicalAddress ini_start = GetInitialProcessBinaryPhysicalAddress();
    const size_t ini_size = GetInitialProcessBinarySize();
    const KPhysicalAddress ini_end = ini_start + ini_size;
    const KPhysicalAddress ini_last = ini_end - 1;
    for (const auto& it : m_system.Kernel().MemoryLayout().GetPhysicalMemoryRegionTree()) {
        if (it.IsDerivedFrom(KMemoryRegionType_DramUserPool)) {
            // Get the manager for the region.
            auto& manager = m_managers[it.GetAttributes()];

            const KPhysicalAddress cur_start = it.GetAddress();
            const KPhysicalAddress cur_last = it.GetLastAddress();
            const KPhysicalAddress cur_end = it.GetEndAddress();

            if (cur_start <= ini_start && ini_last <= cur_last) {
                // Free memory before the ini to the heap.
                if (cur_start != ini_start) {
                    manager.Free(cur_start, (ini_start - cur_start) / PageSize);
                }

                // Open/reserve the ini memory.
                manager.OpenFirst(ini_start, ini_size / PageSize);
                reserved_sizes[it.GetAttributes()] += ini_size;

                // Free memory after the ini to the heap.
                if (ini_last != cur_last) {
                    ASSERT(cur_end != 0);
                    manager.Free(ini_end, (cur_end - ini_end) / PageSize);
                }
            } else {
                // Ensure there's no partial overlap with the ini image.
                if (cur_start <= ini_last) {
                    ASSERT(cur_last < ini_start);
                } else {
                    // Otherwise, check the region for general validity.
                    ASSERT(cur_end != 0);
                }

                // Free the memory to the heap.
                manager.Free(cur_start, it.GetSize() / PageSize);
            }
        }
    }

    // Update the used size for all managers.
    for (size_t i = 0; i < m_num_managers; ++i) {
        m_managers[i].SetInitialUsedHeapSize(reserved_sizes[i]);
    }
}

Result KMemoryManager::InitializeOptimizedMemory(u64 process_id, Pool pool) {
    const u32 pool_index = static_cast<u32>(pool);

    // Lock the pool.
    KScopedLightLock lk(m_pool_locks[pool_index]);

    // Check that we don't already have an optimized process.
    R_UNLESS(!m_has_optimized_process[pool_index], ResultBusy);

    // Set the optimized process id.
    m_optimized_process_ids[pool_index] = process_id;
    m_has_optimized_process[pool_index] = true;

    // Clear the management area for the optimized process.
    for (auto* manager = this->GetFirstManager(pool, Direction::FromFront); manager != nullptr;
         manager = this->GetNextManager(manager, Direction::FromFront)) {
        manager->InitializeOptimizedMemory(m_system.Kernel());
    }

    R_SUCCEED();
}

void KMemoryManager::FinalizeOptimizedMemory(u64 process_id, Pool pool) {
    const u32 pool_index = static_cast<u32>(pool);

    // Lock the pool.
    KScopedLightLock lk(m_pool_locks[pool_index]);

    // If the process was optimized, clear it.
    if (m_has_optimized_process[pool_index] && m_optimized_process_ids[pool_index] == process_id) {
        m_has_optimized_process[pool_index] = false;
    }
}

KPhysicalAddress KMemoryManager::AllocateAndOpenContinuous(size_t num_pages, size_t align_pages,
                                                           u32 option) {
    // Early return if we're allocating no pages.
    if (num_pages == 0) {
        return 0;
    }

    // Lock the pool that we're allocating from.
    const auto [pool, dir] = DecodeOption(option);
    KScopedLightLock lk(m_pool_locks[static_cast<std::size_t>(pool)]);

    // Choose a heap based on our page size request.
    const s32 heap_index = KPageHeap::GetAlignedBlockIndex(num_pages, align_pages);

    // Loop, trying to iterate from each block.
    Impl* chosen_manager = nullptr;
    KPhysicalAddress allocated_block = 0;
    for (chosen_manager = this->GetFirstManager(pool, dir); chosen_manager != nullptr;
         chosen_manager = this->GetNextManager(chosen_manager, dir)) {
        allocated_block = chosen_manager->AllocateAligned(heap_index, num_pages, align_pages);
        if (allocated_block != 0) {
            break;
        }
    }

    // If we failed to allocate, quit now.
    if (allocated_block == 0) {
        return 0;
    }

    // Maintain the optimized memory bitmap, if we should.
    if (m_has_optimized_process[static_cast<size_t>(pool)]) {
        chosen_manager->TrackUnoptimizedAllocation(m_system.Kernel(), allocated_block, num_pages);
    }

    // Open the first reference to the pages.
    chosen_manager->OpenFirst(allocated_block, num_pages);

    return allocated_block;
}

Result KMemoryManager::AllocatePageGroupImpl(KPageGroup* out, size_t num_pages, Pool pool,
                                             Direction dir, bool unoptimized, bool random) {
    // Choose a heap based on our page size request.
    const s32 heap_index = KPageHeap::GetBlockIndex(num_pages);
    R_UNLESS(0 <= heap_index, ResultOutOfMemory);

    // Ensure that we don't leave anything un-freed.
    ON_RESULT_FAILURE {
        for (const auto& it : *out) {
            auto& manager = this->GetManager(it.GetAddress());
            const size_t node_num_pages = std::min<u64>(
                it.GetNumPages(), (manager.GetEndAddress() - it.GetAddress()) / PageSize);
            manager.Free(it.GetAddress(), node_num_pages);
        }
        out->Finalize();
    };

    // Keep allocating until we've allocated all our pages.
    for (s32 index = heap_index; index >= 0 && num_pages > 0; index--) {
        const size_t pages_per_alloc = KPageHeap::GetBlockNumPages(index);
        for (Impl* cur_manager = this->GetFirstManager(pool, dir); cur_manager != nullptr;
             cur_manager = this->GetNextManager(cur_manager, dir)) {
            while (num_pages >= pages_per_alloc) {
                // Allocate a block.
                KPhysicalAddress allocated_block = cur_manager->AllocateBlock(index, random);
                if (allocated_block == 0) {
                    break;
                }

                // Ensure we don't leak the block if we fail.
                ON_RESULT_FAILURE_2 {
                    cur_manager->Free(allocated_block, pages_per_alloc);
                };

                // Add the block to our group.
                R_TRY(out->AddBlock(allocated_block, pages_per_alloc));

                // Maintain the optimized memory bitmap, if we should.
                if (unoptimized) {
                    cur_manager->TrackUnoptimizedAllocation(m_system.Kernel(), allocated_block,
                                                            pages_per_alloc);
                }

                num_pages -= pages_per_alloc;
            }
        }
    }

    // Only succeed if we allocated as many pages as we wanted.
    R_UNLESS(num_pages == 0, ResultOutOfMemory);

    // We succeeded!
    R_SUCCEED();
}

Result KMemoryManager::AllocateAndOpen(KPageGroup* out, size_t num_pages, u32 option) {
    ASSERT(out != nullptr);
    ASSERT(out->GetNumPages() == 0);

    // Early return if we're allocating no pages.
    R_SUCCEED_IF(num_pages == 0);

    // Lock the pool that we're allocating from.
    const auto [pool, dir] = DecodeOption(option);
    KScopedLightLock lk(m_pool_locks[static_cast<size_t>(pool)]);

    // Allocate the page group.
    R_TRY(this->AllocatePageGroupImpl(out, num_pages, pool, dir,
                                      m_has_optimized_process[static_cast<size_t>(pool)], true));

    // Open the first reference to the pages.
    for (const auto& block : *out) {
        KPhysicalAddress cur_address = block.GetAddress();
        size_t remaining_pages = block.GetNumPages();
        while (remaining_pages > 0) {
            // Get the manager for the current address.
            auto& manager = this->GetManager(cur_address);

            // Process part or all of the block.
            const size_t cur_pages =
                std::min(remaining_pages, manager.GetPageOffsetToEnd(cur_address));
            manager.OpenFirst(cur_address, cur_pages);

            // Advance.
            cur_address += cur_pages * PageSize;
            remaining_pages -= cur_pages;
        }
    }

    R_SUCCEED();
}

Result KMemoryManager::AllocateForProcess(KPageGroup* out, size_t num_pages, u32 option,
                                          u64 process_id, u8 fill_pattern) {
    ASSERT(out != nullptr);
    ASSERT(out->GetNumPages() == 0);

    // Decode the option.
    const auto [pool, dir] = DecodeOption(option);

    // Allocate the memory.
    bool optimized;
    {
        // Lock the pool that we're allocating from.
        KScopedLightLock lk(m_pool_locks[static_cast<size_t>(pool)]);

        // Check if we have an optimized process.
        const bool has_optimized = m_has_optimized_process[static_cast<size_t>(pool)];
        const bool is_optimized = m_optimized_process_ids[static_cast<size_t>(pool)] == process_id;

        // Allocate the page group.
        R_TRY(this->AllocatePageGroupImpl(out, num_pages, pool, dir, has_optimized && !is_optimized,
                                          false));

        // Set whether we should optimize.
        optimized = has_optimized && is_optimized;
    }

    // Perform optimized memory tracking, if we should.
    if (optimized) {
        // Iterate over the allocated blocks.
        for (const auto& block : *out) {
            // Get the block extents.
            const KPhysicalAddress block_address = block.GetAddress();
            const size_t block_pages = block.GetNumPages();

            // If it has no pages, we don't need to do anything.
            if (block_pages == 0) {
                continue;
            }

            // Fill all the pages that we need to fill.
            bool any_new = false;
            {
                KPhysicalAddress cur_address = block_address;
                size_t remaining_pages = block_pages;
                while (remaining_pages > 0) {
                    // Get the manager for the current address.
                    auto& manager = this->GetManager(cur_address);

                    // Process part or all of the block.
                    const size_t cur_pages =
                        std::min(remaining_pages, manager.GetPageOffsetToEnd(cur_address));
                    any_new = manager.ProcessOptimizedAllocation(m_system.Kernel(), cur_address,
                                                                 cur_pages, fill_pattern);

                    // Advance.
                    cur_address += cur_pages * PageSize;
                    remaining_pages -= cur_pages;
                }
            }

            // If there are new pages, update tracking for the allocation.
            if (any_new) {
                // Update tracking for the allocation.
                KPhysicalAddress cur_address = block_address;
                size_t remaining_pages = block_pages;
                while (remaining_pages > 0) {
                    // Get the manager for the current address.
                    auto& manager = this->GetManager(cur_address);

                    // Lock the pool for the manager.
                    KScopedLightLock lk(m_pool_locks[static_cast<size_t>(manager.GetPool())]);

                    // Track some or all of the current pages.
                    const size_t cur_pages =
                        std::min(remaining_pages, manager.GetPageOffsetToEnd(cur_address));
                    manager.TrackOptimizedAllocation(m_system.Kernel(), cur_address, cur_pages);

                    // Advance.
                    cur_address += cur_pages * PageSize;
                    remaining_pages -= cur_pages;
                }
            }
        }
    } else {
        // Set all the allocated memory.
        for (const auto& block : *out) {
            m_system.DeviceMemory().buffer.ClearBackingRegion(GetInteger(block.GetAddress()) -
                                                                  Core::DramMemoryMap::Base,
                                                              block.GetSize(), fill_pattern);
        }
    }

    R_SUCCEED();
}

size_t KMemoryManager::Impl::Initialize(KPhysicalAddress address, size_t size,
                                        KVirtualAddress management, KVirtualAddress management_end,
                                        Pool p) {
    // Calculate management sizes.
    const size_t ref_count_size = (size / PageSize) * sizeof(u16);
    const size_t optimize_map_size = CalculateOptimizedProcessOverheadSize(size);
    const size_t manager_size = Common::AlignUp(optimize_map_size + ref_count_size, PageSize);
    const size_t page_heap_size = KPageHeap::CalculateManagementOverheadSize(size);
    const size_t total_management_size = manager_size + page_heap_size;
    ASSERT(manager_size <= total_management_size);
    ASSERT(management + total_management_size <= management_end);
    ASSERT(Common::IsAligned(total_management_size, PageSize));

    // Setup region.
    m_pool = p;
    m_management_region = management;
    m_page_reference_counts.resize(
        Kernel::Board::Nintendo::Nx::KSystemControl::Init::GetIntendedMemorySize() / PageSize);
    ASSERT(Common::IsAligned(GetInteger(m_management_region), PageSize));

    // Initialize the manager's KPageHeap.
    m_heap.Initialize(address, size, management + manager_size, page_heap_size);

    return total_management_size;
}

void KMemoryManager::Impl::InitializeOptimizedMemory(KernelCore& kernel) {
    auto optimize_pa = KPageTable::GetHeapPhysicalAddress(kernel, m_management_region);
    auto* optimize_map = kernel.System().DeviceMemory().GetPointer<u64>(optimize_pa);

    std::memset(optimize_map, 0, CalculateOptimizedProcessOverheadSize(m_heap.GetSize()));
}

void KMemoryManager::Impl::TrackUnoptimizedAllocation(KernelCore& kernel, KPhysicalAddress block,
                                                      size_t num_pages) {
    auto optimize_pa = KPageTable::GetHeapPhysicalAddress(kernel, m_management_region);
    auto* optimize_map = kernel.System().DeviceMemory().GetPointer<u64>(optimize_pa);

    // Get the range we're tracking.
    size_t offset = this->GetPageOffset(block);
    const size_t last = offset + num_pages - 1;

    // Track.
    while (offset <= last) {
        // Mark the page as not being optimized-allocated.
        optimize_map[offset / Common::BitSize<u64>()] &=
            ~(u64(1) << (offset % Common::BitSize<u64>()));

        offset++;
    }
}

void KMemoryManager::Impl::TrackOptimizedAllocation(KernelCore& kernel, KPhysicalAddress block,
                                                    size_t num_pages) {
    auto optimize_pa = KPageTable::GetHeapPhysicalAddress(kernel, m_management_region);
    auto* optimize_map = kernel.System().DeviceMemory().GetPointer<u64>(optimize_pa);

    // Get the range we're tracking.
    size_t offset = this->GetPageOffset(block);
    const size_t last = offset + num_pages - 1;

    // Track.
    while (offset <= last) {
        // Mark the page as being optimized-allocated.
        optimize_map[offset / Common::BitSize<u64>()] |=
            (u64(1) << (offset % Common::BitSize<u64>()));

        offset++;
    }
}

bool KMemoryManager::Impl::ProcessOptimizedAllocation(KernelCore& kernel, KPhysicalAddress block,
                                                      size_t num_pages, u8 fill_pattern) {
    auto& device_memory = kernel.System().DeviceMemory();
    auto optimize_pa = KPageTable::GetHeapPhysicalAddress(kernel, m_management_region);
    auto* optimize_map = device_memory.GetPointer<u64>(optimize_pa);

    // We want to return whether any pages were newly allocated.
    bool any_new = false;

    // Get the range we're processing.
    size_t offset = this->GetPageOffset(block);
    const size_t last = offset + num_pages - 1;

    // Process.
    while (offset <= last) {
        // Check if the page has been optimized-allocated before.
        if ((optimize_map[offset / Common::BitSize<u64>()] &
             (u64(1) << (offset % Common::BitSize<u64>()))) == 0) {
            // If not, it's new.
            any_new = true;

            // Fill the page.
            auto* ptr = device_memory.GetPointer<u8>(m_heap.GetAddress());
            std::memset(ptr + offset * PageSize, fill_pattern, PageSize);
        }

        offset++;
    }

    // Return the number of pages we processed.
    return any_new;
}

size_t KMemoryManager::Impl::CalculateManagementOverheadSize(size_t region_size) {
    const size_t ref_count_size = (region_size / PageSize) * sizeof(u16);
    const size_t optimize_map_size =
        (Common::AlignUp((region_size / PageSize), Common::BitSize<u64>()) /
         Common::BitSize<u64>()) *
        sizeof(u64);
    const size_t manager_meta_size = Common::AlignUp(optimize_map_size + ref_count_size, PageSize);
    const size_t page_heap_size = KPageHeap::CalculateManagementOverheadSize(region_size);
    return manager_meta_size + page_heap_size;
}

} // namespace Kernel