From 95b8c1745a4a9377f629ca204f90add44ac41bf1 Mon Sep 17 00:00:00 2001
From: Ali Mohammad Pur <ali.mpfard@gmail.com>
Date: Thu, 13 Jan 2022 20:59:13 +0330
Subject: [PATCH] AK: Make Variant::visit() prefer overloads accepting T const&
 over T&

This makes the following code behave as expected:

    Variant<int, String> x { some_string() };
    x.visit(
        [](String const&) {}, // Expectation is for this to be called
        [](auto&) {});
---
 AK/Variant.h             | 33 ++++++++++++++++++++++++++++++++-
 Tests/AK/TestVariant.cpp | 14 ++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/AK/Variant.h b/AK/Variant.h
index b872233ff7..175ec3fdad 100644
--- a/AK/Variant.h
+++ b/AK/Variant.h
@@ -80,13 +80,41 @@ struct Variant<IndexType, InitialIndex> {
 
 template<typename IndexType, typename... Ts>
 struct VisitImpl {
+    template<typename RT, typename T, size_t I, typename Fn>
+    static consteval u64 get_explicitly_named_overload_if_exists()
+    {
+        // If we're not allowed to make a member function pointer and call it directly (without explicitly resolving it),
+        // we have a templated function on our hands (or a function overload set).
+        // in such cases, we don't have an explicitly named overload, and we would have to select it.
+        if constexpr (requires { (declval<Fn>().*(&Fn::operator()))(declval<T>()); })
+            return 1ull << I;
+
+        return 0;
+    }
+
+    template<typename ReturnType, typename T, typename Visitor, auto... Is>
+    static consteval bool should_invoke_const_overload(IndexSequence<Is...>)
+    {
+        // Scan over all the different visitor functions, if none of them are suitable for calling with `T const&`, avoid calling that first.
+        return ((get_explicitly_named_overload_if_exists<ReturnType, T, Is, typename Visitor::Types::template Type<Is>>()) | ...) != 0;
+    }
+
     template<typename Self, typename Visitor, IndexType CurrentIndex = 0>
     ALWAYS_INLINE static constexpr decltype(auto) visit(Self& self, IndexType id, const void* data, Visitor&& visitor) requires(CurrentIndex < sizeof...(Ts))
     {
         using T = typename TypeList<Ts...>::template Type<CurrentIndex>;
 
-        if (id == CurrentIndex)
+        if (id == CurrentIndex) {
+            // Check if Visitor::operator() is an explicitly typed function (as opposed to a templated function)
+            // if so, try to call that with `T const&` first before copying the Variant's const-ness.
+            // This emulates normal C++ call semantics where templated functions are considered last, after all non-templated overloads
+            // are checked and found to be unusable.
+            using ReturnType = decltype(visitor(*bit_cast<T*>(data)));
+            if constexpr (should_invoke_const_overload<ReturnType, T, Visitor>(MakeIndexSequence<Visitor::Types::size>()))
+                return visitor(*bit_cast<AddConst<T>*>(data));
+
             return visitor(*bit_cast<CopyConst<Self, T>*>(data));
+        }
 
         if constexpr ((CurrentIndex + 1) < sizeof...(Ts))
             return visit<Self, Visitor, CurrentIndex + 1>(self, id, data, forward<Visitor>(visitor));
@@ -424,6 +452,9 @@ private:
 
     template<typename... Fs>
     struct Visitor : Fs... {
+        using Types = TypeList<Fs...>;
+        static_assert(Types::size < 64, "Variant::visit() can take a maximum of 64 visit functions.");
+
         Visitor(Fs&&... args)
             : Fs(forward<Fs>(args))...
         {
diff --git a/Tests/AK/TestVariant.cpp b/Tests/AK/TestVariant.cpp
index af87669671..f55d176537 100644
--- a/Tests/AK/TestVariant.cpp
+++ b/Tests/AK/TestVariant.cpp
@@ -48,6 +48,20 @@ TEST_CASE(visit_const)
         [&](auto const&) {});
 
     EXPECT(correct);
+
+    correct = false;
+    auto the_value_but_not_const = the_value;
+    the_value_but_not_const.visit(
+        [&](String const&) { correct = true; },
+        [&](auto&) {});
+
+    EXPECT(correct);
+
+    correct = false;
+    the_value_but_not_const.visit(
+        [&]<typename T>(T&) { correct = !IsConst<T>; });
+
+    EXPECT(correct);
 }
 
 TEST_CASE(destructor)