6 files changed, 128 insertions, 26 deletions

diff --git a/src/core/hle/kernel/k_memory_block.h b/src/core/hle/kernel/k_memory_block.h
index 8437cb659..ef3f61321 100644
--- a/src/core/hle/kernel/k_memory_block.h
+++ b/src/core/hle/kernel/k_memory_block.h
@@ -192,7 +192,7 @@ enum class KMemoryAttribute : u8 {
     Uncached = static_cast<u8>(Svc::MemoryAttribute::Uncached),
     PermissionLocked = static_cast<u8>(Svc::MemoryAttribute::PermissionLocked),
 
-    SetMask = Uncached,
+    SetMask = Uncached | PermissionLocked,
 };
 DECLARE_ENUM_FLAG_OPERATORS(KMemoryAttribute);
 
@@ -339,6 +339,10 @@ public:
         return this->GetEndAddress() - 1;
     }
 
+    constexpr KMemoryState GetState() const {
+        return m_memory_state;
+    }
+
     constexpr u16 GetIpcLockCount() const {
         return m_ipc_lock_count;
     }
@@ -456,6 +460,13 @@ public:
         }
     }
 
+    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));
diff --git a/src/core/hle/kernel/k_memory_block_manager.cpp b/src/core/hle/kernel/k_memory_block_manager.cpp
index ab75f550e..58a1e7216 100644
--- a/src/core/hle/kernel/k_memory_block_manager.cpp
+++ b/src/core/hle/kernel/k_memory_block_manager.cpp
@@ -160,8 +160,8 @@ void KMemoryBlockManager::Update(KMemoryBlockManagerUpdateAllocator* allocator,
             }
 
             // Update block state.
-            it->Update(state, perm, attr, cur_address == address, static_cast<u8>(set_disable_attr),
-                       static_cast<u8>(clear_disable_attr));
+            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();
         }
@@ -175,7 +175,9 @@ void KMemoryBlockManager::UpdateIfMatch(KMemoryBlockManagerUpdateAllocator* allo
                                         KProcessAddress address, size_t num_pages,
                                         KMemoryState test_state, KMemoryPermission test_perm,
                                         KMemoryAttribute test_attr, KMemoryState state,
-                                        KMemoryPermission perm, KMemoryAttribute attr) {
+                                        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));
@@ -214,7 +216,8 @@ void KMemoryBlockManager::UpdateIfMatch(KMemoryBlockManagerUpdateAllocator* allo
             }
 
             // Update block state.
-            it->Update(state, perm, attr, false, 0, 0);
+            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 {
@@ -284,6 +287,65 @@ void KMemoryBlockManager::UpdateLock(KMemoryBlockManagerUpdateAllocator* allocat
     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.
diff --git a/src/core/hle/kernel/k_memory_block_manager.h b/src/core/hle/kernel/k_memory_block_manager.h
index 96496e990..cb7b6f430 100644
--- a/src/core/hle/kernel/k_memory_block_manager.h
+++ b/src/core/hle/kernel/k_memory_block_manager.h
@@ -115,7 +115,11 @@ public:
     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);
+                       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(
diff --git a/src/core/hle/kernel/k_page_table.cpp b/src/core/hle/kernel/k_page_table.cpp
index ebc540316..387f2c962 100644
--- a/src/core/hle/kernel/k_page_table.cpp
+++ b/src/core/hle/kernel/k_page_table.cpp
@@ -505,7 +505,7 @@ Result KPageTable::UnmapCodeMemory(KProcessAddress dst_address, KProcessAddress
     R_TRY(this->CheckMemoryStateContiguous(
         std::addressof(num_dst_allocator_blocks), dst_address, size, KMemoryState::FlagCanCodeAlias,
         KMemoryState::FlagCanCodeAlias, KMemoryPermission::None, KMemoryPermission::None,
-        KMemoryAttribute::All, KMemoryAttribute::None));
+        KMemoryAttribute::All & ~KMemoryAttribute::PermissionLocked, KMemoryAttribute::None));
 
     // Determine whether any pages being unmapped are code.
     bool any_code_pages = false;
@@ -1770,7 +1770,11 @@ Result KPageTable::MapPhysicalMemory(KProcessAddress address, size_t size) {
                 m_memory_block_manager.UpdateIfMatch(
                     std::addressof(allocator), address, size / PageSize, KMemoryState::Free,
                     KMemoryPermission::None, KMemoryAttribute::None, KMemoryState::Normal,
-                    KMemoryPermission::UserReadWrite, KMemoryAttribute::None);
+                    KMemoryPermission::UserReadWrite, KMemoryAttribute::None,
+                    address == this->GetAliasRegionStart()
+                        ? KMemoryBlockDisableMergeAttribute::Normal
+                        : KMemoryBlockDisableMergeAttribute::None,
+                    KMemoryBlockDisableMergeAttribute::None);
 
                 R_SUCCEED();
             }
@@ -1868,6 +1872,13 @@ Result KPageTable::UnmapPhysicalMemory(KProcessAddress address, size_t size) {
 
     // Iterate over the memory, unmapping as we go.
     auto it = m_memory_block_manager.FindIterator(cur_address);
+
+    const auto clear_merge_attr =
+        (it->GetState() == KMemoryState::Normal &&
+         it->GetAddress() == this->GetAliasRegionStart() && it->GetAddress() == address)
+            ? KMemoryBlockDisableMergeAttribute::Normal
+            : KMemoryBlockDisableMergeAttribute::None;
+
     while (true) {
         // Check that the iterator is valid.
         ASSERT(it != m_memory_block_manager.end());
@@ -1905,7 +1916,7 @@ Result KPageTable::UnmapPhysicalMemory(KProcessAddress address, size_t size) {
     m_memory_block_manager.Update(std::addressof(allocator), address, size / PageSize,
                                   KMemoryState::Free, KMemoryPermission::None,
                                   KMemoryAttribute::None, KMemoryBlockDisableMergeAttribute::None,
-                                  KMemoryBlockDisableMergeAttribute::None);
+                                  clear_merge_attr);
 
     // We succeeded.
     R_SUCCEED();
@@ -2650,11 +2661,18 @@ Result KPageTable::SetMemoryAttribute(KProcessAddress addr, size_t size, u32 mas
     size_t num_allocator_blocks;
     constexpr auto AttributeTestMask =
         ~(KMemoryAttribute::SetMask | KMemoryAttribute::DeviceShared);
-    R_TRY(this->CheckMemoryState(
-        std::addressof(old_state), std::addressof(old_perm), std::addressof(old_attr),
-        std::addressof(num_allocator_blocks), addr, size, KMemoryState::FlagCanChangeAttribute,
-        KMemoryState::FlagCanChangeAttribute, KMemoryPermission::None, KMemoryPermission::None,
-        AttributeTestMask, KMemoryAttribute::None, ~AttributeTestMask));
+    const KMemoryState state_test_mask =
+        static_cast<KMemoryState>(((mask & static_cast<u32>(KMemoryAttribute::Uncached))
+                                       ? static_cast<u32>(KMemoryState::FlagCanChangeAttribute)
+                                       : 0) |
+                                  ((mask & static_cast<u32>(KMemoryAttribute::PermissionLocked))
+                                       ? static_cast<u32>(KMemoryState::FlagCanPermissionLock)
+                                       : 0));
+    R_TRY(this->CheckMemoryState(std::addressof(old_state), std::addressof(old_perm),
+                                 std::addressof(old_attr), std::addressof(num_allocator_blocks),
+                                 addr, size, state_test_mask, state_test_mask,
+                                 KMemoryPermission::None, KMemoryPermission::None,
+                                 AttributeTestMask, KMemoryAttribute::None, ~AttributeTestMask));
 
     // Create an update allocator.
     Result allocator_result{ResultSuccess};
@@ -2662,18 +2680,17 @@ Result KPageTable::SetMemoryAttribute(KProcessAddress addr, size_t size, u32 mas
                                                  m_memory_block_slab_manager, num_allocator_blocks);
     R_TRY(allocator_result);
 
-    // Determine the new attribute.
-    const KMemoryAttribute new_attr =
-        static_cast<KMemoryAttribute>(((old_attr & static_cast<KMemoryAttribute>(~mask)) |
-                                       static_cast<KMemoryAttribute>(attr & mask)));
-
-    // Perform operation.
-    this->Operate(addr, num_pages, old_perm, OperationType::ChangePermissionsAndRefresh);
+    // If we need to, perform a change attribute operation.
+    if (True(KMemoryAttribute::Uncached & static_cast<KMemoryAttribute>(mask))) {
+        // Perform operation.
+        R_TRY(this->Operate(addr, num_pages, old_perm,
+                            OperationType::ChangePermissionsAndRefreshAndFlush, 0));
+    }
 
     // Update the blocks.
-    m_memory_block_manager.Update(std::addressof(allocator), addr, num_pages, old_state, old_perm,
-                                  new_attr, KMemoryBlockDisableMergeAttribute::None,
-                                  KMemoryBlockDisableMergeAttribute::None);
+    m_memory_block_manager.UpdateAttribute(std::addressof(allocator), addr, num_pages,
+                                           static_cast<KMemoryAttribute>(mask),
+                                           static_cast<KMemoryAttribute>(attr));
 
     R_SUCCEED();
 }
@@ -3086,6 +3103,7 @@ Result KPageTable::Operate(KProcessAddress addr, size_t num_pages, KMemoryPermis
     }
     case OperationType::ChangePermissions:
     case OperationType::ChangePermissionsAndRefresh:
+    case OperationType::ChangePermissionsAndRefreshAndFlush:
         break;
     default:
         ASSERT(false);
diff --git a/src/core/hle/kernel/k_page_table.h b/src/core/hle/kernel/k_page_table.h
index e69498f02..cbcbb6f62 100644
--- a/src/core/hle/kernel/k_page_table.h
+++ b/src/core/hle/kernel/k_page_table.h
@@ -222,7 +222,8 @@ private:
         Unmap = 3,
         ChangePermissions = 4,
         ChangePermissionsAndRefresh = 5,
-        Separate = 6,
+        ChangePermissionsAndRefreshAndFlush = 6,
+        Separate = 7,
     };
 
     static constexpr KMemoryAttribute DefaultMemoryIgnoreAttr =
diff --git a/src/core/hle/kernel/svc/svc_memory.cpp b/src/core/hle/kernel/svc/svc_memory.cpp
index 2cab74127..372684094 100644
--- a/src/core/hle/kernel/svc/svc_memory.cpp
+++ b/src/core/hle/kernel/svc/svc_memory.cpp
@@ -108,10 +108,16 @@ Result SetMemoryAttribute(Core::System& system, u64 address, u64 size, u32 mask,
     R_UNLESS((address < address + size), ResultInvalidCurrentMemory);
 
     // Validate the attribute and mask.
-    constexpr u32 SupportedMask = static_cast<u32>(MemoryAttribute::Uncached);
+    constexpr u32 SupportedMask =
+        static_cast<u32>(MemoryAttribute::Uncached | MemoryAttribute::PermissionLocked);
     R_UNLESS((mask | attr) == mask, ResultInvalidCombination);
     R_UNLESS((mask | attr | SupportedMask) == SupportedMask, ResultInvalidCombination);
 
+    // Check that permission locked is either being set or not masked.
+    R_UNLESS((static_cast<Svc::MemoryAttribute>(mask) & Svc::MemoryAttribute::PermissionLocked) ==
+                 (static_cast<Svc::MemoryAttribute>(attr) & Svc::MemoryAttribute::PermissionLocked),
+             ResultInvalidCombination);
+
     // Validate that the region is in range for the current process.
     auto& page_table{GetCurrentProcess(system.Kernel()).GetPageTable()};
     R_UNLESS(page_table.Contains(address, size), ResultInvalidCurrentMemory);