6 files changed, 141 insertions, 83 deletions

diff --git a/src/core/hle/kernel/k_memory_block.h b/src/core/hle/kernel/k_memory_block.h
index fd491146f..9e51c33ce 100644
--- a/src/core/hle/kernel/k_memory_block.h
+++ b/src/core/hle/kernel/k_memory_block.h
@@ -120,7 +120,7 @@ static_assert(static_cast<u32>(KMemoryState::CodeOut) == 0x00402015);
 
 enum class KMemoryPermission : u8 {
     None = 0,
-    Mask = static_cast<u8>(~None),
+    All = static_cast<u8>(~None),
 
     Read = 1 << 0,
     Write = 1 << 1,
diff --git a/src/core/hle/kernel/k_page_table.cpp b/src/core/hle/kernel/k_page_table.cpp
index f2f88c147..4da509224 100644
--- a/src/core/hle/kernel/k_page_table.cpp
+++ b/src/core/hle/kernel/k_page_table.cpp
@@ -264,9 +264,9 @@ ResultCode KPageTable::InitializeForProcess(FileSys::ProgramAddressSpaceType as_
     ASSERT(heap_last < stack_start || stack_last < heap_start);
     ASSERT(heap_last < kmap_start || kmap_last < heap_start);
 
-    current_heap_addr = heap_region_start;
-    heap_capacity = 0;
-    physical_memory_usage = 0;
+    current_heap_end = heap_region_start;
+    max_heap_size = 0;
+    mapped_physical_memory_size = 0;
     memory_pool = pool;
 
     page_table_impl.Resize(address_space_width, PageBits);
@@ -306,7 +306,7 @@ ResultCode KPageTable::MapProcessCodeMemory(VAddr dst_addr, VAddr src_addr, std:
     KMemoryState state{};
     KMemoryPermission perm{};
     CASCADE_CODE(CheckMemoryState(&state, &perm, nullptr, src_addr, size, KMemoryState::All,
-                                  KMemoryState::Normal, KMemoryPermission::Mask,
+                                  KMemoryState::Normal, KMemoryPermission::All,
                                   KMemoryPermission::ReadAndWrite, KMemoryAttribute::Mask,
                                   KMemoryAttribute::None, KMemoryAttribute::IpcAndDeviceMapped));
 
@@ -465,7 +465,7 @@ ResultCode KPageTable::MapPhysicalMemory(VAddr addr, std::size_t size) {
 
     MapPhysicalMemory(page_linked_list, addr, end_addr);
 
-    physical_memory_usage += remaining_size;
+    mapped_physical_memory_size += remaining_size;
 
     const std::size_t num_pages{size / PageSize};
     block_manager->Update(addr, num_pages, KMemoryState::Free, KMemoryPermission::None,
@@ -507,7 +507,7 @@ ResultCode KPageTable::UnmapPhysicalMemory(VAddr addr, std::size_t size) {
 
     auto process{system.Kernel().CurrentProcess()};
     process->GetResourceLimit()->Release(LimitableResource::PhysicalMemory, mapped_size);
-    physical_memory_usage -= mapped_size;
+    mapped_physical_memory_size -= mapped_size;
 
     return ResultSuccess;
 }
@@ -554,7 +554,7 @@ ResultCode KPageTable::Map(VAddr dst_addr, VAddr src_addr, std::size_t size) {
     KMemoryState src_state{};
     CASCADE_CODE(CheckMemoryState(
         &src_state, nullptr, nullptr, src_addr, size, KMemoryState::FlagCanAlias,
-        KMemoryState::FlagCanAlias, KMemoryPermission::Mask, KMemoryPermission::ReadAndWrite,
+        KMemoryState::FlagCanAlias, KMemoryPermission::All, KMemoryPermission::ReadAndWrite,
         KMemoryAttribute::Mask, KMemoryAttribute::None, KMemoryAttribute::IpcAndDeviceMapped));
 
     if (IsRegionMapped(dst_addr, size)) {
@@ -593,7 +593,7 @@ ResultCode KPageTable::Unmap(VAddr dst_addr, VAddr src_addr, std::size_t size) {
     KMemoryState src_state{};
     CASCADE_CODE(CheckMemoryState(
         &src_state, nullptr, nullptr, src_addr, size, KMemoryState::FlagCanAlias,
-        KMemoryState::FlagCanAlias, KMemoryPermission::Mask, KMemoryPermission::None,
+        KMemoryState::FlagCanAlias, KMemoryPermission::All, KMemoryPermission::None,
         KMemoryAttribute::Mask, KMemoryAttribute::Locked, KMemoryAttribute::IpcAndDeviceMapped));
 
     KMemoryPermission dst_perm{};
@@ -784,7 +784,7 @@ ResultCode KPageTable::ReserveTransferMemory(VAddr addr, std::size_t size, KMemo
     CASCADE_CODE(CheckMemoryState(
         &state, nullptr, &attribute, addr, size,
         KMemoryState::FlagCanTransfer | KMemoryState::FlagReferenceCounted,
-        KMemoryState::FlagCanTransfer | KMemoryState::FlagReferenceCounted, KMemoryPermission::Mask,
+        KMemoryState::FlagCanTransfer | KMemoryState::FlagReferenceCounted, KMemoryPermission::All,
         KMemoryPermission::ReadAndWrite, KMemoryAttribute::Mask, KMemoryAttribute::None,
         KMemoryAttribute::IpcAndDeviceMapped));
 
@@ -859,61 +859,125 @@ ResultCode KPageTable::SetMemoryAttribute(VAddr addr, std::size_t size, KMemoryA
     return ResultSuccess;
 }
 
-ResultCode KPageTable::SetHeapCapacity(std::size_t new_heap_capacity) {
+ResultCode KPageTable::SetMaxHeapSize(std::size_t size) {
+    // Lock the table.
     std::lock_guard lock{page_table_lock};
-    heap_capacity = new_heap_capacity;
-    return ResultSuccess;
-}
 
-ResultVal<VAddr> KPageTable::SetHeapSize(std::size_t size) {
+    // Only process page tables are allowed to set heap size.
+    ASSERT(!this->IsKernel());
 
-    if (size > heap_region_end - heap_region_start) {
-        return ResultOutOfMemory;
-    }
+    max_heap_size = size;
 
-    const u64 previous_heap_size{GetHeapSize()};
-
-    UNIMPLEMENTED_IF_MSG(previous_heap_size > size, "Heap shrink is unimplemented");
+    return ResultSuccess;
+}
 
-    // Increase the heap size
+ResultCode KPageTable::SetHeapSize(VAddr* out, std::size_t size) {
+    // Try to perform a reduction in heap, instead of an extension.
+    VAddr cur_address{};
+    std::size_t allocation_size{};
     {
-        std::lock_guard lock{page_table_lock};
-
-        const u64 delta{size - previous_heap_size};
-
-        // Reserve memory for the heap extension.
-        KScopedResourceReservation memory_reservation(
-            system.Kernel().CurrentProcess()->GetResourceLimit(), LimitableResource::PhysicalMemory,
-            delta);
-
-        if (!memory_reservation.Succeeded()) {
-            LOG_ERROR(Kernel, "Could not reserve heap extension of size {:X} bytes", delta);
-            return ResultLimitReached;
+        // Lock the table.
+        std::lock_guard lk(page_table_lock);
+
+        // Validate that setting heap size is possible at all.
+        R_UNLESS(!is_kernel, ResultOutOfMemory);
+        R_UNLESS(size <= static_cast<std::size_t>(heap_region_end - heap_region_start),
+                 ResultOutOfMemory);
+        R_UNLESS(size <= max_heap_size, ResultOutOfMemory);
+
+        if (size < GetHeapSize()) {
+            // The size being requested is less than the current size, so we need to free the end of
+            // the heap.
+
+            // Validate memory state.
+            std::size_t num_allocator_blocks;
+            R_TRY(this->CheckMemoryState(std::addressof(num_allocator_blocks),
+                                         heap_region_start + size, GetHeapSize() - size,
+                                         KMemoryState::All, KMemoryState::Normal,
+                                         KMemoryPermission::All, KMemoryPermission::ReadAndWrite,
+                                         KMemoryAttribute::All, KMemoryAttribute::None));
+
+            // Unmap the end of the heap.
+            const auto num_pages = (GetHeapSize() - size) / PageSize;
+            R_TRY(Operate(heap_region_start + size, num_pages, KMemoryPermission::None,
+                          OperationType::Unmap));
+
+            // Release the memory from the resource limit.
+            system.Kernel().CurrentProcess()->GetResourceLimit()->Release(
+                LimitableResource::PhysicalMemory, num_pages * PageSize);
+
+            // Apply the memory block update.
+            block_manager->Update(heap_region_start + size, num_pages, KMemoryState::Free,
+                                  KMemoryPermission::None, KMemoryAttribute::None);
+
+            // Update the current heap end.
+            current_heap_end = heap_region_start + size;
+
+            // Set the output.
+            *out = heap_region_start;
+            return ResultSuccess;
+        } else if (size == GetHeapSize()) {
+            // The size requested is exactly the current size.
+            *out = heap_region_start;
+            return ResultSuccess;
+        } else {
+            // We have to allocate memory. Determine how much to allocate and where while the table
+            // is locked.
+            cur_address = current_heap_end;
+            allocation_size = size - GetHeapSize();
         }
+    }
 
-        KPageLinkedList page_linked_list;
-        const std::size_t num_pages{delta / PageSize};
+    // Reserve memory for the heap extension.
+    KScopedResourceReservation memory_reservation(
+        system.Kernel().CurrentProcess()->GetResourceLimit(), LimitableResource::PhysicalMemory,
+        allocation_size);
+    R_UNLESS(memory_reservation.Succeeded(), ResultLimitReached);
 
-        CASCADE_CODE(
-            system.Kernel().MemoryManager().Allocate(page_linked_list, num_pages, memory_pool));
+    // Allocate pages for the heap extension.
+    KPageLinkedList page_linked_list;
+    R_TRY(system.Kernel().MemoryManager().Allocate(page_linked_list, allocation_size / PageSize,
+                                                   memory_pool));
 
-        if (IsRegionMapped(current_heap_addr, delta)) {
-            return ResultInvalidCurrentMemory;
+    // Map the pages.
+    {
+        // Lock the table.
+        std::lock_guard lk(page_table_lock);
+
+        // Ensure that the heap hasn't changed since we began executing.
+        ASSERT(cur_address == current_heap_end);
+
+        // Check the memory state.
+        std::size_t num_allocator_blocks{};
+        R_TRY(this->CheckMemoryState(std::addressof(num_allocator_blocks), current_heap_end,
+                                     allocation_size, KMemoryState::All, KMemoryState::Free,
+                                     KMemoryPermission::None, KMemoryPermission::None,
+                                     KMemoryAttribute::None, KMemoryAttribute::None));
+
+        // Map the pages.
+        const auto num_pages = allocation_size / PageSize;
+        R_TRY(Operate(current_heap_end, num_pages, page_linked_list, OperationType::MapGroup));
+
+        // Clear all the newly allocated pages.
+        for (std::size_t cur_page = 0; cur_page < num_pages; ++cur_page) {
+            std::memset(system.Memory().GetPointer(current_heap_end + (cur_page * PageSize)), 0,
+                        PageSize);
         }
 
-        CASCADE_CODE(
-            Operate(current_heap_addr, num_pages, page_linked_list, OperationType::MapGroup));
-
-        // Succeeded in allocation, commit the resource reservation
+        // We succeeded, so commit our memory reservation.
         memory_reservation.Commit();
 
-        block_manager->Update(current_heap_addr, num_pages, KMemoryState::Normal,
-                              KMemoryPermission::ReadAndWrite);
+        // Apply the memory block update.
+        block_manager->Update(current_heap_end, num_pages, KMemoryState::Normal,
+                              KMemoryPermission::ReadAndWrite, KMemoryAttribute::None);
 
-        current_heap_addr = heap_region_start + size;
-    }
+        // Update the current heap end.
+        current_heap_end = heap_region_start + size;
 
-    return heap_region_start;
+        // Set the output.
+        *out = heap_region_start;
+        return ResultSuccess;
+    }
 }
 
 ResultVal<VAddr> KPageTable::AllocateAndMapMemory(std::size_t needed_num_pages, std::size_t align,
@@ -1005,7 +1069,7 @@ ResultCode KPageTable::LockForCodeMemory(VAddr addr, std::size_t size) {
 
     if (const ResultCode result{CheckMemoryState(
             nullptr, &old_perm, nullptr, addr, size, KMemoryState::FlagCanCodeMemory,
-            KMemoryState::FlagCanCodeMemory, KMemoryPermission::Mask,
+            KMemoryState::FlagCanCodeMemory, KMemoryPermission::All,
             KMemoryPermission::UserReadWrite, KMemoryAttribute::All, KMemoryAttribute::None)};
         result.IsError()) {
         return result;
@@ -1058,9 +1122,8 @@ ResultCode KPageTable::InitializeMemoryLayout(VAddr start, VAddr end) {
 
 bool KPageTable::IsRegionMapped(VAddr address, u64 size) {
     return CheckMemoryState(address, size, KMemoryState::All, KMemoryState::Free,
-                            KMemoryPermission::Mask, KMemoryPermission::None,
-                            KMemoryAttribute::Mask, KMemoryAttribute::None,
-                            KMemoryAttribute::IpcAndDeviceMapped)
+                            KMemoryPermission::All, KMemoryPermission::None, KMemoryAttribute::Mask,
+                            KMemoryAttribute::None, KMemoryAttribute::IpcAndDeviceMapped)
         .IsError();
 }
 
diff --git a/src/core/hle/kernel/k_page_table.h b/src/core/hle/kernel/k_page_table.h
index db08ea8ce..564410dca 100644
--- a/src/core/hle/kernel/k_page_table.h
+++ b/src/core/hle/kernel/k_page_table.h
@@ -50,8 +50,8 @@ public:
     ResultCode SetMemoryPermission(VAddr addr, std::size_t size, Svc::MemoryPermission perm);
     ResultCode SetMemoryAttribute(VAddr addr, std::size_t size, KMemoryAttribute mask,
                                   KMemoryAttribute value);
-    ResultCode SetHeapCapacity(std::size_t new_heap_capacity);
-    ResultVal<VAddr> SetHeapSize(std::size_t size);
+    ResultCode SetMaxHeapSize(std::size_t size);
+    ResultCode SetHeapSize(VAddr* out, std::size_t size);
     ResultVal<VAddr> AllocateAndMapMemory(std::size_t needed_num_pages, std::size_t align,
                                           bool is_map_only, VAddr region_start,
                                           std::size_t region_num_pages, KMemoryState state,
@@ -183,14 +183,15 @@ public:
     constexpr VAddr GetAliasCodeRegionSize() const {
         return alias_code_region_end - alias_code_region_start;
     }
+    size_t GetNormalMemorySize() {
+        std::lock_guard lk(page_table_lock);
+        return GetHeapSize() + mapped_physical_memory_size;
+    }
     constexpr std::size_t GetAddressSpaceWidth() const {
         return address_space_width;
     }
-    constexpr std::size_t GetHeapSize() {
-        return current_heap_addr - heap_region_start;
-    }
-    constexpr std::size_t GetTotalHeapSize() {
-        return GetHeapSize() + physical_memory_usage;
+    constexpr std::size_t GetHeapSize() const {
+        return current_heap_end - heap_region_start;
     }
     constexpr bool IsInsideAddressSpace(VAddr address, std::size_t size) const {
         return address_space_start <= address && address + size - 1 <= address_space_end - 1;
@@ -270,10 +271,8 @@ private:
     VAddr code_region_end{};
     VAddr alias_code_region_start{};
     VAddr alias_code_region_end{};
-    VAddr current_heap_addr{};
 
-    std::size_t heap_capacity{};
-    std::size_t physical_memory_usage{};
+    std::size_t mapped_physical_memory_size{};
     std::size_t max_heap_size{};
     std::size_t max_physical_memory_size{};
     std::size_t address_space_width{};
diff --git a/src/core/hle/kernel/k_process.cpp b/src/core/hle/kernel/k_process.cpp
index aee313995..73f8bc4fe 100644
--- a/src/core/hle/kernel/k_process.cpp
+++ b/src/core/hle/kernel/k_process.cpp
@@ -172,7 +172,7 @@ void KProcess::DecrementThreadCount() {
 
 u64 KProcess::GetTotalPhysicalMemoryAvailable() const {
     const u64 capacity{resource_limit->GetFreeValue(LimitableResource::PhysicalMemory) +
-                       page_table->GetTotalHeapSize() + GetSystemResourceSize() + image_size +
+                       page_table->GetNormalMemorySize() + GetSystemResourceSize() + image_size +
                        main_thread_stack_size};
     if (const auto pool_size = kernel.MemoryManager().GetSize(KMemoryManager::Pool::Application);
         capacity != pool_size) {
@@ -189,7 +189,7 @@ u64 KProcess::GetTotalPhysicalMemoryAvailableWithoutSystemResource() const {
 }
 
 u64 KProcess::GetTotalPhysicalMemoryUsed() const {
-    return image_size + main_thread_stack_size + page_table->GetTotalHeapSize() +
+    return image_size + main_thread_stack_size + page_table->GetNormalMemorySize() +
            GetSystemResourceSize();
 }
 
@@ -410,8 +410,8 @@ void KProcess::Run(s32 main_thread_priority, u64 stack_size) {
     resource_limit->Reserve(LimitableResource::Threads, 1);
     resource_limit->Reserve(LimitableResource::PhysicalMemory, main_thread_stack_size);
 
-    const std::size_t heap_capacity{memory_usage_capacity - main_thread_stack_size - image_size};
-    ASSERT(!page_table->SetHeapCapacity(heap_capacity).IsError());
+    const std::size_t heap_capacity{memory_usage_capacity - (main_thread_stack_size + image_size)};
+    ASSERT(!page_table->SetMaxHeapSize(heap_capacity).IsError());
 
     ChangeStatus(ProcessStatus::Running);
 
diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index 68cb47211..63e2dff19 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -135,24 +135,15 @@ enum class ResourceLimitValueType {
 } // Anonymous namespace
 
 /// Set the process heap to a given Size. It can both extend and shrink the heap.
-static ResultCode SetHeapSize(Core::System& system, VAddr* heap_addr, u64 heap_size) {
-    LOG_TRACE(Kernel_SVC, "called, heap_size=0x{:X}", heap_size);
+static ResultCode SetHeapSize(Core::System& system, VAddr* out_address, u64 size) {
+    LOG_TRACE(Kernel_SVC, "called, heap_size=0x{:X}", size);
 
-    // Size must be a multiple of 0x200000 (2MB) and be equal to or less than 8GB.
-    if ((heap_size % 0x200000) != 0) {
-        LOG_ERROR(Kernel_SVC, "The heap size is not a multiple of 2MB, heap_size=0x{:016X}",
-                  heap_size);
-        return ResultInvalidSize;
-    }
-
-    if (heap_size >= 0x200000000) {
-        LOG_ERROR(Kernel_SVC, "The heap size is not less than 8GB, heap_size=0x{:016X}", heap_size);
-        return ResultInvalidSize;
-    }
-
-    auto& page_table{system.Kernel().CurrentProcess()->PageTable()};
+    // Validate size.
+    R_UNLESS(Common::IsAligned(size, HeapSizeAlignment), ResultInvalidSize);
+    R_UNLESS(size < MainMemorySizeMax, ResultInvalidSize);
 
-    CASCADE_RESULT(*heap_addr, page_table.SetHeapSize(heap_size));
+    // Set the heap size.
+    R_TRY(system.Kernel().CurrentProcess()->PageTable().SetHeapSize(out_address, size));
 
     return ResultSuccess;
 }
diff --git a/src/core/hle/kernel/svc_common.h b/src/core/hle/kernel/svc_common.h
index 60ea2c405..25de6e437 100644
--- a/src/core/hle/kernel/svc_common.h
+++ b/src/core/hle/kernel/svc_common.h
@@ -5,6 +5,7 @@
 #pragma once
 
 #include "common/common_types.h"
+#include "common/literals.h"
 
 namespace Kernel {
 using Handle = u32;
@@ -12,9 +13,13 @@ using Handle = u32;
 
 namespace Kernel::Svc {
 
+using namespace Common::Literals;
+
 constexpr s32 ArgumentHandleCountMax = 0x40;
 constexpr u32 HandleWaitMask{1u << 30};
 
+constexpr inline std::size_t HeapSizeAlignment = 2_MiB;
+
 constexpr inline Handle InvalidHandle = Handle(0);
 
 enum PseudoHandle : Handle {