From 45a0ba2167172b15e3f5b124d3e3c48a0f84b2ba Mon Sep 17 00:00:00 2001
From: Dan Klishch <danilklishch@gmail.com>
Date: Sun, 10 Mar 2024 20:52:50 -0400
Subject: [PATCH] AK: Introduce AK::enumerate

Co-Authored-By: Tim Flynn <trflynn89@pm.me>
---
 AK/Enumerate.h             | 65 ++++++++++++++++++++++++++++++++++++++
 Tests/AK/CMakeLists.txt    |  1 +
 Tests/AK/TestEnumerate.cpp | 51 ++++++++++++++++++++++++++++++
 3 files changed, 117 insertions(+)
 create mode 100644 AK/Enumerate.h
 create mode 100644 Tests/AK/TestEnumerate.cpp

diff --git a/AK/Enumerate.h b/AK/Enumerate.h
new file mode 100644
index 00000000000..d824605e7dd
--- /dev/null
+++ b/AK/Enumerate.h
@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) 2024, Tim Flynn <trflynn89@serenityos.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/StdLibExtras.h>
+
+namespace AK {
+
+namespace Detail {
+template<typename Iterable>
+class Enumerator {
+    using IteratorType = decltype(declval<Iterable>().begin());
+    using ValueType = decltype(*declval<IteratorType>());
+
+    struct Enumeration {
+        size_t index { 0 };
+        ValueType value;
+    };
+
+public:
+    Enumerator(Iterable&& iterable)
+        : m_iterable(forward<Iterable>(iterable))
+        , m_iterator(m_iterable.begin())
+        , m_end(m_iterable.end())
+    {
+    }
+
+    Enumerator const& begin() const { return *this; }
+    Enumerator const& end() const { return *this; }
+
+    Enumeration operator*() { return { m_index, *m_iterator }; }
+    Enumeration operator*() const { return { m_index, *m_iterator }; }
+
+    bool operator!=(Enumerator const&) const { return m_iterator != m_end; }
+
+    void operator++()
+    {
+        ++m_index;
+        ++m_iterator;
+    }
+
+private:
+    Iterable m_iterable;
+
+    size_t m_index { 0 };
+    IteratorType m_iterator;
+    IteratorType const m_end;
+};
+}
+
+template<typename T>
+auto enumerate(T&& range)
+{
+    return Detail::Enumerator<T> { forward<T>(range) };
+}
+
+}
+
+#ifdef USING_AK_GLOBALLY
+using AK::enumerate;
+#endif
diff --git a/Tests/AK/CMakeLists.txt b/Tests/AK/CMakeLists.txt
index 5346d366e89..ec85150dae9 100644
--- a/Tests/AK/CMakeLists.txt
+++ b/Tests/AK/CMakeLists.txt
@@ -25,6 +25,7 @@ set(AK_TEST_SOURCES
     TestDoublyLinkedList.cpp
     TestEndian.cpp
     TestEnumBits.cpp
+    TestEnumerate.cpp
     TestFind.cpp
     TestFixedArray.cpp
     TestFixedPoint.cpp
diff --git a/Tests/AK/TestEnumerate.cpp b/Tests/AK/TestEnumerate.cpp
new file mode 100644
index 00000000000..dffc2869b62
--- /dev/null
+++ b/Tests/AK/TestEnumerate.cpp
@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2024, Dan Klishch <danilklishch@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <LibTest/TestCase.h>
+
+#include <AK/Enumerate.h>
+#include <AK/Span.h>
+#include <AK/Vector.h>
+
+struct IndexAndValue {
+    size_t index;
+    int value;
+
+    bool operator==(IndexAndValue const&) const = default;
+};
+
+TEST_CASE(enumerate)
+{
+    {
+        Vector<IndexAndValue> result;
+        for (auto [i, value] : enumerate(Vector { 1, 2, 3, 4 })) {
+            result.append({ i, value });
+        }
+        EXPECT_EQ(result, (Vector<IndexAndValue> { { 0, 1 }, { 1, 2 }, { 2, 3 }, { 3, 4 } }));
+    }
+
+    {
+        Vector<IndexAndValue> result;
+        Vector<int> values = { 9, 8, 7, 6 };
+        for (auto [i, value] : enumerate(values)) {
+            static_assert(SameAs<decltype(value), int&>);
+            result.append({ i, value });
+            value = static_cast<int>(i);
+        }
+        EXPECT_EQ(result, (Vector<IndexAndValue> { { 0, 9 }, { 1, 8 }, { 2, 7 }, { 3, 6 } }));
+        EXPECT_EQ(values, (Vector<int> { 0, 1, 2, 3 }));
+    }
+
+    {
+        Vector<IndexAndValue> result;
+        Vector<int> const& values = { 9, 8, 7, 6 };
+        for (auto [i, value] : enumerate(values)) {
+            static_assert(SameAs<decltype(value), int const&>);
+            result.append({ i, value });
+        }
+        EXPECT_EQ(result, (Vector<IndexAndValue> { { 0, 9 }, { 1, 8 }, { 2, 7 }, { 3, 6 } }));
+    }
+}