From 8a01167c7dc5acbbee636c66017f5f35e8766865 Mon Sep 17 00:00:00 2001
From: Itamar <itamar8910@gmail.com>
Date: Sat, 8 May 2021 12:27:12 +0300
Subject: [PATCH] AK: Add missing GenericTraits<NonnullRefPtr>

This enables us to use keys of type NonnullRefPtr in HashMaps and
HashTables.

This commit also includes fixes in various places that used
HashMap<T, NonnullRefPtr<U>>::get() and expected to get an
Optional<NonnullRefPtr<U>> and now get an Optional<U*>.
---
 AK/NonnullRefPtr.h                                       | 9 +++++++++
 .../LanguageServers/Cpp/ParserAutoComplete.cpp           | 1 +
 Userland/DevTools/HackStudio/LanguageServers/FileDB.cpp  | 3 ++-
 Userland/Libraries/LibELF/DynamicLinker.cpp              | 8 +++++---
 Userland/Services/Taskbar/main.cpp                       | 2 +-
 Userland/Services/WindowServer/ClientConnection.h        | 4 ++--
 6 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/AK/NonnullRefPtr.h b/AK/NonnullRefPtr.h
index b4ea20b451..ac80672bc8 100644
--- a/AK/NonnullRefPtr.h
+++ b/AK/NonnullRefPtr.h
@@ -9,6 +9,7 @@
 #include <AK/Assertions.h>
 #include <AK/Atomic.h>
 #include <AK/Format.h>
+#include <AK/Traits.h>
 #include <AK/Types.h>
 #ifdef KERNEL
 #    include <Kernel/Arch/x86/CPU.h>
@@ -335,5 +336,13 @@ inline void swap(NonnullRefPtr<T>& a, NonnullRefPtr<U>& b)
 
 }
 
+template<typename T>
+struct Traits<NonnullRefPtr<T>> : public GenericTraits<NonnullRefPtr<T>> {
+    using PeekType = T*;
+    using ConstPeekType = const T*;
+    static unsigned hash(const NonnullRefPtr<T>& p) { return ptr_hash(p.ptr()); }
+    static bool equals(const NonnullRefPtr<T>& a, const NonnullRefPtr<T>& b) { return a.ptr() == b.ptr(); }
+};
+
 using AK::adopt_ref;
 using AK::NonnullRefPtr;
diff --git a/Userland/DevTools/HackStudio/LanguageServers/Cpp/ParserAutoComplete.cpp b/Userland/DevTools/HackStudio/LanguageServers/Cpp/ParserAutoComplete.cpp
index 2aedb992a1..5a7df291a8 100644
--- a/Userland/DevTools/HackStudio/LanguageServers/Cpp/ParserAutoComplete.cpp
+++ b/Userland/DevTools/HackStudio/LanguageServers/Cpp/ParserAutoComplete.cpp
@@ -7,6 +7,7 @@
 #include "ParserAutoComplete.h"
 #include <AK/Assertions.h>
 #include <AK/HashTable.h>
+#include <AK/OwnPtr.h>
 #include <LibCpp/AST.h>
 #include <LibCpp/Lexer.h>
 #include <LibCpp/Parser.h>
diff --git a/Userland/DevTools/HackStudio/LanguageServers/FileDB.cpp b/Userland/DevTools/HackStudio/LanguageServers/FileDB.cpp
index 6b3e7c25a8..474bbc13b1 100644
--- a/Userland/DevTools/HackStudio/LanguageServers/FileDB.cpp
+++ b/Userland/DevTools/HackStudio/LanguageServers/FileDB.cpp
@@ -7,6 +7,7 @@
 #include "FileDB.h"
 
 #include <AK/LexicalPath.h>
+#include <AK/NonnullRefPtr.h>
 #include <LibCore/File.h>
 
 namespace LanguageServers {
@@ -18,7 +19,7 @@ RefPtr<const GUI::TextDocument> FileDB::get(const String& filename) const
     if (!document_optional.has_value())
         return nullptr;
 
-    return document_optional.value();
+    return adopt_ref(*document_optional.value());
 }
 
 RefPtr<GUI::TextDocument> FileDB::get(const String& filename)
diff --git a/Userland/Libraries/LibELF/DynamicLinker.cpp b/Userland/Libraries/LibELF/DynamicLinker.cpp
index 3423e517a4..6551d9db55 100644
--- a/Userland/Libraries/LibELF/DynamicLinker.cpp
+++ b/Userland/Libraries/LibELF/DynamicLinker.cpp
@@ -311,7 +311,7 @@ static Result<NonnullRefPtr<DynamicLoader>, DlErrorMessage> load_main_library(co
         loader.load_stage_4();
     }
 
-    return main_library_loader;
+    return NonnullRefPtr<DynamicLoader>(*main_library_loader);
 }
 
 static Result<void, DlErrorMessage> __dlclose(void* handle)
@@ -377,7 +377,8 @@ static Result<void*, DlErrorMessage> __dlopen(const char* filename, int flags)
     auto existing_elf_object = s_global_objects.get(library_name);
     if (existing_elf_object.has_value()) {
         // It's up to the caller to release the ref with dlclose().
-        return &existing_elf_object->leak_ref();
+        existing_elf_object.value()->ref();
+        return *existing_elf_object;
     }
 
     VERIFY(!library_name.is_empty());
@@ -406,7 +407,8 @@ static Result<void*, DlErrorMessage> __dlopen(const char* filename, int flags)
         return DlErrorMessage { "Could not load ELF object." };
 
     // It's up to the caller to release the ref with dlclose().
-    return &object->leak_ref();
+    object.value()->ref();
+    return *object;
 }
 
 static Result<void*, DlErrorMessage> __dlsym(void* handle, const char* symbol_name)
diff --git a/Userland/Services/Taskbar/main.cpp b/Userland/Services/Taskbar/main.cpp
index 9d2041ef3e..557621dc68 100644
--- a/Userland/Services/Taskbar/main.cpp
+++ b/Userland/Services/Taskbar/main.cpp
@@ -144,7 +144,7 @@ NonnullRefPtr<GUI::Menu> build_system_menu()
                 dbgln("App {} has icon with size {}", app.name, icon->size());
         }
 
-        auto parent_menu = app_category_menus.get(app.category).value_or(*system_menu);
+        auto parent_menu = app_category_menus.get(app.category).value_or(system_menu.ptr());
         parent_menu->add_action(GUI::Action::create(app.name, icon, [app_identifier](auto&) {
             dbgln("Activated app with ID {}", app_identifier);
             const auto& bin = g_apps[app_identifier].executable;
diff --git a/Userland/Services/WindowServer/ClientConnection.h b/Userland/Services/WindowServer/ClientConnection.h
index 0fba35da8d..7be8fd5401 100644
--- a/Userland/Services/WindowServer/ClientConnection.h
+++ b/Userland/Services/WindowServer/ClientConnection.h
@@ -46,14 +46,14 @@ public:
         auto menu = m_menus.get(menu_id);
         if (!menu.has_value())
             return nullptr;
-        return const_cast<Menu*>(menu.value().ptr());
+        return menu.value();
     }
     const Menu* find_menu_by_id(int menu_id) const
     {
         auto menu = m_menus.get(menu_id);
         if (!menu.has_value())
             return nullptr;
-        return menu.value().ptr();
+        return menu.value();
     }
 
     template<typename Callback>