From b6cb561c47a42ecc954a1f80c624f6b08d556cb8 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Wed, 29 Jul 2020 02:29:46 -0700
Subject: [PATCH] kern: fix slab heap atomics

---
 .../arch/arm64/kern_k_slab_heap_impl.hpp      | 69 +++++++++++++++++++
 .../include/mesosphere/kern_k_slab_heap.hpp   | 34 +++++----
 2 files changed, 88 insertions(+), 15 deletions(-)
 create mode 100644 libraries/libmesosphere/include/mesosphere/arch/arm64/kern_k_slab_heap_impl.hpp

diff --git a/libraries/libmesosphere/include/mesosphere/arch/arm64/kern_k_slab_heap_impl.hpp b/libraries/libmesosphere/include/mesosphere/arch/arm64/kern_k_slab_heap_impl.hpp
new file mode 100644
index 000000000..c52f28528
--- /dev/null
+++ b/libraries/libmesosphere/include/mesosphere/arch/arm64/kern_k_slab_heap_impl.hpp
@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) 2018-2020 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#pragma once
+#include <vapours.hpp>
+#include <mesosphere/kern_select_cpu.hpp>
+
+namespace ams::kern::arch::arm64 {
+
+    template<typename T>
+    concept SlabHeapNode = requires (T &t) {
+        { t.next } -> std::convertible_to<T *>;
+    };
+
+    template<typename T> requires SlabHeapNode<T>
+    ALWAYS_INLINE T *AllocateFromSlabAtomic(T **head) {
+        u32 tmp;
+        T *node, *next;
+
+        __asm__ __volatile__(
+            "1:\n"
+            "    ldaxr  %[node], [%[head]]\n"
+            "    cbz    %[node], 2f\n"
+            "    ldr    %[next], [%[node]]\n"
+            "    stlxr  %w[tmp], %[next], [%[head]]\n"
+            "    cbnz   %w[tmp], 1b\n"
+            "    b      3f\n"
+            "2:\n"
+            "    clrex\n"
+            "3:\n"
+            : [tmp]"=&r"(tmp), [node]"=&r"(node), [next]"=&r"(next), [head]"+&r"(head)
+            :
+            : "cc", "memory"
+        );
+
+        return node;
+    }
+
+    template<typename T> requires SlabHeapNode<T>
+    ALWAYS_INLINE void FreeToSlabAtomic(T **head, T *node) {
+        u32 tmp;
+        T *next;
+
+        __asm__ __volatile__(
+            "1:\n"
+            "    ldaxr  %[next], [%[head]]\n"
+            "    str    %[next], [%[node]]\n"
+            "    stlxr  %w[tmp], %[node], [%[head]]\n"
+            "    cbnz   %w[tmp], 1b\n"
+            "2:\n"
+            : [tmp]"=&r"(tmp), [node]"+&r"(node), [next]"=&r"(next), [head]"+&r"(head)
+            :
+            : "cc", "memory"
+        );
+    }
+
+}
diff --git a/libraries/libmesosphere/include/mesosphere/kern_k_slab_heap.hpp b/libraries/libmesosphere/include/mesosphere/kern_k_slab_heap.hpp
index 2d24df367..494777700 100644
--- a/libraries/libmesosphere/include/mesosphere/kern_k_slab_heap.hpp
+++ b/libraries/libmesosphere/include/mesosphere/kern_k_slab_heap.hpp
@@ -17,6 +17,21 @@
 #include <mesosphere/kern_common.hpp>
 #include <mesosphere/kern_k_typed_address.hpp>
 
+#if defined(ATMOSPHERE_ARCH_ARM64)
+
+    #include <mesosphere/arch/arm64/kern_k_slab_heap_impl.hpp>
+    namespace ams::kern {
+        using ams::kern::arch::arm64::AllocateFromSlabAtomic;
+        using ams::kern::arch::arm64::FreeToSlabAtomic;
+    }
+
+#else
+
+    #error "Unknown architecture for KSlabHeapImpl"
+
+#endif
+
+
 namespace ams::kern {
 
     namespace impl {
@@ -29,8 +44,8 @@ namespace ams::kern {
                     Node *next;
                 };
             private:
-                std::atomic<Node *> head;
-                size_t              obj_size;
+                Node * head;
+                size_t obj_size;
             public:
                 constexpr KSlabHeapImpl() : head(nullptr), obj_size(0) { MESOSPHERE_ASSERT_THIS(); }
 
@@ -50,15 +65,7 @@ namespace ams::kern {
                 void *Allocate() {
                     MESOSPHERE_ASSERT_THIS();
 
-                    Node *ret = this->head.load();
-
-                    do {
-                        if (AMS_UNLIKELY(ret == nullptr)) {
-                            break;
-                        }
-                    } while (!this->head.compare_exchange_weak(ret, ret->next));
-
-                    return ret;
+                    return AllocateFromSlabAtomic(std::addressof(this->head));
                 }
 
                 void Free(void *obj) {
@@ -66,10 +73,7 @@ namespace ams::kern {
 
                     Node *node = reinterpret_cast<Node *>(obj);
 
-                    Node *cur_head = this->head.load();
-                    do {
-                        node->next = cur_head;
-                    } while (!this->head.compare_exchange_weak(cur_head, node));
+                    return FreeToSlabAtomic(std::addressof(this->head), node);
                 }
         };