From 7345a952cac2949ddb66b2a8b271159d7d1237e1 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Tue, 2 Jul 2024 20:51:28 +0200
Subject: [PATCH] libluna: Implement hash table iteration

---
 libluna/include/luna/HashMap.h   | 10 +++++
 libluna/include/luna/HashTable.h | 75 ++++++++++++++++++++++++++++++++
 tests/libluna/TestHashTable.cpp  | 19 ++++++++
 3 files changed, 104 insertions(+)

diff --git a/libluna/include/luna/HashMap.h b/libluna/include/luna/HashMap.h
index e2573b01..83672e15 100644
--- a/libluna/include/luna/HashMap.h
+++ b/libluna/include/luna/HashMap.h
@@ -154,6 +154,16 @@ template <typename K, typename V> struct HashMap
         m_table.clear();
     }
 
+    const HashTable<HashPair<K, V>>::HashTableIterator begin()
+    {
+        return m_table.begin();
+    }
+
+    const HashTable<HashPair<K, V>>::HashTableIterator end()
+    {
+        return m_table.end();
+    }
+
   private:
     HashTable<HashPair<K, V>> m_table;
 };
diff --git a/libluna/include/luna/HashTable.h b/libluna/include/luna/HashTable.h
index a29e86d7..53ba041f 100644
--- a/libluna/include/luna/HashTable.h
+++ b/libluna/include/luna/HashTable.h
@@ -12,6 +12,8 @@
 #include <luna/Heap.h>
 #include <luna/Option.h>
 
+template <typename K, typename V> class HashMap;
+
 /**
  * @brief A table of values with best-case constant time lookup.
  *
@@ -181,6 +183,77 @@ template <typename T> class HashTable
         clear();
     }
 
+    struct HashTableIterator
+    {
+      private:
+        HashTableIterator(usize index, HashTable<T>& table) : m_table(table)
+        {
+            if (index != -1)
+            {
+                m_index = index;
+                return;
+            }
+
+            for (usize i = 0; i < m_table.m_capacity; i++)
+            {
+                auto& opt = m_table.m_buckets[i];
+                if (opt.has_value())
+                {
+                    m_index = i;
+                    return;
+                }
+            }
+
+            m_index = m_table.m_capacity;
+        }
+
+        mutable usize m_index;
+        HashTable<T>& m_table;
+
+      public:
+        T& operator*()
+        {
+            return *m_table.m_buckets[m_index].value_ptr();
+        }
+
+        const T& operator*() const
+        {
+            return *m_table.m_buckets[m_index].value_ptr();
+        }
+
+        void operator++() const
+        {
+            for (usize i = m_index + 1; i < m_table.m_capacity; i++)
+            {
+                auto& opt = m_table.m_buckets[i];
+                if (opt.has_value())
+                {
+                    m_index = i;
+                    return;
+                }
+            }
+
+            m_index = m_table.m_capacity;
+        }
+
+        bool operator!=(HashTableIterator& other) const
+        {
+            return m_index != other.m_index || &m_table != &other.m_table;
+        }
+
+        friend class HashTable<T>;
+    };
+
+    const HashTableIterator begin()
+    {
+        return { (usize)-1, *this };
+    }
+
+    const HashTableIterator end()
+    {
+        return { m_capacity, *this };
+    }
+
   private:
     bool should_grow()
     {
@@ -223,4 +296,6 @@ template <typename T> class HashTable
     usize m_size { 0 };
     // FIXME: Randomize this to protect against hash table attacks.
     u64 m_salt { 0 };
+
+    friend class HashTableIterator;
 };
diff --git a/tests/libluna/TestHashTable.cpp b/tests/libluna/TestHashTable.cpp
index 5b410a24..06b64191 100644
--- a/tests/libluna/TestHashTable.cpp
+++ b/tests/libluna/TestHashTable.cpp
@@ -80,6 +80,24 @@ TestResult test_hash_map_set_and_get()
     test_success;
 }
 
+TestResult test_hash_table_iteration()
+{
+    HashTable<int> table;
+
+    validate(TRY(table.try_set(1)));
+    validate(TRY(table.try_set(2)));
+    validate(TRY(table.try_set(3)));
+    validate(TRY(table.try_set(4)));
+
+    int sum = 0;
+
+    for (const int& value : table) { sum += value; }
+
+    validate(sum == 10);
+
+    test_success;
+}
+
 Result<void> test_main()
 {
     test_prelude;
@@ -89,6 +107,7 @@ Result<void> test_main()
     run_test(test_hash_table_remove);
     run_test(test_hash_table_duplicates);
     run_test(test_hash_map_set_and_get);
+    run_test(test_hash_table_iteration);
 
     return {};
 }