FIX: Fixing ARM64 execution when booting from bootloader.

FIX: Now works both for ARM64 and AMD64 EFI backends.
TODO: The ACB backend needs to be done too.

Signed-off-by: Amlal El Mahrouss <amlal.elmahrouss@icloud.com>

7 files changed, 81 insertions, 3 deletions

diff --git a/dev/ZKAKit/KernelKit/IDLLObject.h b/dev/ZKAKit/KernelKit/IDLLObject.h
index bd1aa5d0..e14e335f 100644
--- a/dev/ZKAKit/KernelKit/IDLLObject.h
+++ b/dev/ZKAKit/KernelKit/IDLLObject.h
@@ -12,6 +12,8 @@
 #include <NewKit/Defines.h>
 #include <CompilerKit/CompilerKit.h>
 
+#define ZKA_DLL_OBJECT : public IDLLObject
+
 namespace Kernel
 {
 	/// @brief DLL class object. A handle to a shared library.
diff --git a/dev/ZKAKit/KernelKit/IPEFDLLObject.h b/dev/ZKAKit/KernelKit/IPEFDLLObject.h
index 9447a98f..03dcd53f 100644
--- a/dev/ZKAKit/KernelKit/IPEFDLLObject.h
+++ b/dev/ZKAKit/KernelKit/IPEFDLLObject.h
@@ -22,7 +22,7 @@ namespace Kernel
 	 * @brief Shared Library class
 	 * Load library from this class
 	 */
-	class IPEFDLLObject final : public IDLLObject
+	class IPEFDLLObject final ZKA_DLL_OBJECT
 	{
 	public:
 		explicit IPEFDLLObject() = default;
diff --git a/dev/ZKAKit/KernelKit/PEFCodeMgr.h b/dev/ZKAKit/KernelKit/PEFCodeMgr.h
index 5474e4b6..c42701ce 100644
--- a/dev/ZKAKit/KernelKit/PEFCodeMgr.h
+++ b/dev/ZKAKit/KernelKit/PEFCodeMgr.h
@@ -12,6 +12,10 @@
 #include <NewKit/KString.h>
 #include <KernelKit/FileMgr.h>
 
+#ifndef INC_PROCESS_SCHEDULER_H
+#include <KernelKit/UserProcessScheduler.h>
+#endif
+
 #define kPefApplicationMime "application/vnd-zka-executable"
 
 namespace Kernel
diff --git a/dev/ZKAKit/KernelKit/ThreadLocalStorage.h b/dev/ZKAKit/KernelKit/ThreadLocalStorage.h
index 93e0a9a4..2c0a8810 100644
--- a/dev/ZKAKit/KernelKit/ThreadLocalStorage.h
+++ b/dev/ZKAKit/KernelKit/ThreadLocalStorage.h
@@ -8,6 +8,7 @@
 #define KERNELKIT_TLS_H
 
 #include <NewKit/Defines.h>
+#include <NewKit/ErrorOr.h>
 
 ///! @brief Thread Local Storage for minoskrnl.
 
@@ -38,6 +39,16 @@ T* tls_new_ptr(void) noexcept;
 template <typename T>
 Kernel::Boolean tls_delete_ptr(T* ptr) noexcept;
 
+//! @brief Delete process pointer.
+//! @param obj The pointer to delete.
+template <typename T>
+inline Kernel::Bool tls_delete_ptr(Kernel::ErrorOr<T> obj) noexcept;
+
+//! @brief Delete process pointer.
+//! @param obj The pointer to delete.
+template <typename T>
+inline Kernel::Bool tls_delete_ptr(Kernel::ErrorOr<T*> obj) noexcept;
+
 template <typename T, typename... Args>
 T* tls_new_class(Args&&... args);
 
diff --git a/dev/ZKAKit/KernelKit/ThreadLocalStorage.inl b/dev/ZKAKit/KernelKit/ThreadLocalStorage.inl
index d9850c01..3683bc4b 100644
--- a/dev/ZKAKit/KernelKit/ThreadLocalStorage.inl
+++ b/dev/ZKAKit/KernelKit/ThreadLocalStorage.inl
@@ -7,6 +7,7 @@
 //! @file ThreadLocalStorage.inl
 //! @brief Allocate resources from the process's heap storage.
 
+#include "NewKit/ErrorOr.h"
 #ifndef INC_PROCESS_SCHEDULER_H
 #include <KernelKit/UserProcessScheduler.h>
 #endif
@@ -40,7 +41,25 @@ inline Kernel::Bool tls_delete_ptr(T* obj) noexcept
 	auto ref_process = UserProcessScheduler::The().GetCurrentProcess();
 	MUST_PASS(ref_process);
 
-	return ref_process.Leak().Delete(obj, sizeof(T));
+	ErrorOr<T*> obj_wrapped{obj};
+
+	return ref_process.Leak().Delete(obj_wrapped, sizeof(T));
+}
+
+//! @brief Delete process pointer.
+//! @param obj The pointer to delete.
+template <typename T>
+inline Kernel::Bool tls_delete_ptr(Kernel::ErrorOr<T> obj) noexcept
+{
+	return tls_delete_ptr(obj.Leak());
+}
+
+//! @brief Delete process pointer.
+//! @param obj The pointer to delete.
+template <typename T>
+inline Kernel::Bool tls_delete_ptr(Kernel::ErrorOr<T*> obj) noexcept
+{
+	return tls_delete_ptr(obj->Leak());
 }
 
 /// @brief Allocate a C++ class, and then call the constructor of it.
diff --git a/dev/ZKAKit/KernelKit/UserProcessScheduler.h b/dev/ZKAKit/KernelKit/UserProcessScheduler.h
index a3a6c8bd..c1559085 100644
--- a/dev/ZKAKit/KernelKit/UserProcessScheduler.h
+++ b/dev/ZKAKit/KernelKit/UserProcessScheduler.h
@@ -219,7 +219,8 @@ namespace Kernel
 		///! @brief TLS free.
 		///! @param ptr the pointer to free.
 		///! @param sz the size of it.
-		Boolean Delete(ErrorOr<VoidPtr> ptr, const SizeT& sz);
+		template <typename T>
+		Boolean Delete(ErrorOr<T*> ptr, const SizeT& sz);
 
 		///! @brief Wakes up threads.
 		Void Wake(const Bool wakeup = false);
@@ -324,6 +325,7 @@ namespace Kernel
 } // namespace Kernel
 
 #include <KernelKit/ThreadLocalStorage.h>
+#include <KernelKit/UserProcessScheduler.inl>
 
 ////////////////////////////////////////////////////
 
diff --git a/dev/ZKAKit/KernelKit/UserProcessScheduler.inl b/dev/ZKAKit/KernelKit/UserProcessScheduler.inl
new file mode 100644
index 00000000..b99fafbf
--- /dev/null
+++ b/dev/ZKAKit/KernelKit/UserProcessScheduler.inl
@@ -0,0 +1,40 @@
+namespace Kernel
+{
+	/***********************************************************************************/
+	/** @brief Free pointer from usage. */
+	/***********************************************************************************/
+
+    template <typename T>
+	Boolean UserProcess::Delete(ErrorOr<T*> ptr, const SizeT& sz)
+	{
+		if (!ptr ||
+			sz == 0)
+			return No;
+
+		UserProcessHeapList* entry = this->MemoryHeap;
+
+		while (entry != nullptr)
+		{
+			if (entry->MemoryEntry == ptr.Leak().Leak())
+			{
+#ifdef __ZKA_AMD64__
+				auto pd = hal_read_cr3();
+				hal_write_cr3(this->VMRegister);
+
+				auto ret = mm_delete_heap(entry->MemoryEntry);
+
+				hal_write_cr3(pd);
+
+				return ret;
+#else
+				Bool ret = mm_delete_heap(ptr.Leak().Leak());
+				return ret;
+#endif
+			}
+
+			entry = entry->MemoryNext;
+		}
+
+		return No;
+	}
+} // namespace Kernel
\ No newline at end of file