From 26a4fe7bfd33321625be9c9b8a8920e3a74fe4ec Mon Sep 17 00:00:00 2001
From: AntoinePrv <AntoinePrv@users.noreply.github.com>
Date: Wed, 31 May 2023 17:40:19 +0200
Subject: [PATCH] Add Python-like set operations to flat_set

---
 libmamba/include/mamba/util/flat_set.hpp   | 233 ++++++++++++++++++++-
 libmamba/src/solv-cpp/queue.cpp            |   5 +
 libmamba/src/solv-cpp/queue.hpp            |   2 +
 libmamba/tests/src/solv-cpp/test_queue.cpp |   7 +
 libmamba/tests/src/util/test_flat_set.cpp  | 102 +++++++++
 5 files changed, 342 insertions(+), 7 deletions(-)

diff --git a/libmamba/include/mamba/util/flat_set.hpp b/libmamba/include/mamba/util/flat_set.hpp
index 443a1413e9..ba8fb14096 100644
--- a/libmamba/include/mamba/util/flat_set.hpp
+++ b/libmamba/include/mamba/util/flat_set.hpp
@@ -30,6 +30,7 @@ namespace mamba::util
     public:
 
         using Base = std::vector<Key, Allocator>;
+        using Self = flat_set<Key, Compare, Allocator>;
         using typename Base::allocator_type;
         using typename Base::const_iterator;
         using typename Base::const_reverse_iterator;
@@ -69,7 +70,6 @@ namespace mamba::util
         flat_set& operator=(const flat_set&) = default;
         flat_set& operator=(flat_set&&) = default;
 
-        bool contains(const value_type&) const;
         const value_type& front() const noexcept;
         const value_type& back() const noexcept;
 
@@ -91,6 +91,16 @@ namespace mamba::util
         const_iterator erase(const_iterator first, const_iterator last);
         size_type erase(const value_type& value);
 
+        bool contains(const value_type&) const;
+        bool is_disjoint_of(const Self& other) const;
+        bool is_subset_of(const Self& other) const;
+        bool is_superset_of(const Self& other) const;
+
+        static Self union_(const Self& lhs, const Self& rhs);
+        static Self intersection(const Self& lhs, const Self& rhs);
+        static Self difference(const Self& lhs, const Self& rhs);
+        static Self symetric_difference(const Self& lhs, const Self& rhs);
+
     private:
 
         key_compare m_compare;
@@ -131,6 +141,46 @@ namespace mamba::util
     bool
     operator!=(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
 
+    /** Return whether the first set is a subset of the second. */
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator<=(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
+    /** Return whether the first set is a strict subset of the second. */
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator<(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
+    /** Return whether the first set is a superset of the second. */
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator>=(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
+    /** Return whether the first set is a strict superset of the second. */
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator>(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
+    /** Compute the set union. */
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator|(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
+    /** Compute the set intersection. */
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator&(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
+    /** Compute the set difference. */
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator-(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
+    /** Compute the set symetric difference. */
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator^(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs);
+
     /*******************************
      *  vector_set Implementation  *
      *******************************/
@@ -177,12 +227,6 @@ namespace mamba::util
         sort_and_remove_duplicates();
     }
 
-    template <typename K, typename C, typename A>
-    auto flat_set<K, C, A>::contains(const value_type& value) const -> bool
-    {
-        return std::binary_search(begin(), end(), value);
-    }
-
     template <typename K, typename C, typename A>
     auto flat_set<K, C, A>::front() const noexcept -> const value_type&
     {
@@ -292,6 +336,51 @@ namespace mamba::util
         return 1;
     }
 
+    template <typename K, typename C, typename A>
+    auto flat_set<K, C, A>::contains(const value_type& value) const -> bool
+    {
+        return std::binary_search(begin(), end(), value);
+    }
+
+    namespace detail
+    {
+        /**
+         * Check if two sorted range have an empty intersection.
+         *
+         * Edited from https://en.cppreference.com/w/cpp/algorithm/set_intersection
+         * Distributed under the terms of the Copyright/CC-BY-SA License.
+         * The full license can be found at the address
+         * https://en.cppreference.com/w/Cppreference:Copyright/CC-BY-SA
+         */
+        template <class InputIt1, class InputIt2, class Compare>
+        bool
+        set_disjoint(InputIt1 first1, InputIt1 last1, InputIt2 first2, InputIt2 last2, Compare comp)
+        {
+            while (first1 != last1 && first2 != last2)
+            {
+                if (comp(*first1, *first2))
+                {
+                    ++first1;
+                }
+                else
+                {
+                    if (!comp(*first2, *first1))
+                    {
+                        return false;  // *first1 and *first2 are equivalent.
+                    }
+                    ++first2;
+                }
+            }
+            return true;
+        }
+    }
+
+    template <typename K, typename C, typename A>
+    auto flat_set<K, C, A>::is_disjoint_of(const Self& other) const -> bool
+    {
+        return detail::set_disjoint(cbegin(), cend(), other.cbegin(), other.cend(), m_compare);
+    }
+
     template <typename K, typename C, typename A>
     bool operator==(const flat_set<K, C, A>& lhs, const flat_set<K, C, A>& rhs)
     {
@@ -305,5 +394,135 @@ namespace mamba::util
         return !(lhs == rhs);
     }
 
+    template <typename K, typename C, typename A>
+    auto flat_set<K, C, A>::is_subset_of(const Self& other) const -> bool
+    {
+        return std::includes(other.cbegin(), other.cend(), cbegin(), cend(), m_compare);
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator<=(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return lhs.is_subset_of(rhs);
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator<(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return (lhs.size() < rhs.size()) && (lhs <= rhs);
+    }
+
+    template <typename K, typename C, typename A>
+    auto flat_set<K, C, A>::is_superset_of(const Self& other) const -> bool
+    {
+        return other.is_subset_of(*this);
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator>=(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return lhs.is_superset_of(rhs);
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    bool
+    operator>(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return rhs < lhs;
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    auto flat_set<Key, Compare, Allocator>::union_(const Self& lhs, const Self& rhs) -> Self
+    {
+        auto out = flat_set<Key, Compare, Allocator>();
+        out.reserve(std::max(lhs.size(), rhs.size()));  // lower bound
+        std::set_union(
+            lhs.cbegin(),
+            lhs.cend(),
+            rhs.cbegin(),
+            rhs.cend(),
+            std::back_inserter(static_cast<Base&>(out)),
+            lhs.m_compare
+        );
+        return out;
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator|(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return flat_set<Key, Compare, Allocator>::union_(lhs, rhs);
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    auto flat_set<Key, Compare, Allocator>::intersection(const Self& lhs, const Self& rhs) -> Self
+    {
+        auto out = flat_set<Key, Compare, Allocator>();
+        std::set_intersection(
+            lhs.cbegin(),
+            lhs.cend(),
+            rhs.cbegin(),
+            rhs.cend(),
+            std::back_inserter(static_cast<Base&>(out)),
+            lhs.m_compare
+        );
+        return out;
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator&(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return flat_set<Key, Compare, Allocator>::intersection(lhs, rhs);
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    auto flat_set<Key, Compare, Allocator>::difference(const Self& lhs, const Self& rhs) -> Self
+    {
+        auto out = flat_set<Key, Compare, Allocator>();
+        std::set_difference(
+            lhs.cbegin(),
+            lhs.cend(),
+            rhs.cbegin(),
+            rhs.cend(),
+            std::back_inserter(static_cast<Base&>(out)),
+            lhs.m_compare
+        );
+        return out;
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator-(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return flat_set<Key, Compare, Allocator>::difference(lhs, rhs);
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    auto flat_set<Key, Compare, Allocator>::symetric_difference(const Self& lhs, const Self& rhs)
+        -> Self
+    {
+        auto out = flat_set<Key, Compare, Allocator>();
+        std::set_symmetric_difference(
+            lhs.cbegin(),
+            lhs.cend(),
+            rhs.cbegin(),
+            rhs.cend(),
+            std::back_inserter(static_cast<Base&>(out)),
+            lhs.m_compare
+        );
+        return out;
+    }
+
+    template <typename Key, typename Compare, typename Allocator>
+    flat_set<Key, Compare, Allocator>
+    operator^(const flat_set<Key, Compare, Allocator>& lhs, const flat_set<Key, Compare, Allocator>& rhs)
+    {
+        return flat_set<Key, Compare, Allocator>::symetric_difference(lhs, rhs);
+    }
+
 }
 #endif
diff --git a/libmamba/src/solv-cpp/queue.cpp b/libmamba/src/solv-cpp/queue.cpp
index b19b9e6816..84582ca731 100644
--- a/libmamba/src/solv-cpp/queue.cpp
+++ b/libmamba/src/solv-cpp/queue.cpp
@@ -250,6 +250,11 @@ namespace mamba::solv
         return m_queue.elements;
     }
 
+    auto ObjQueue::contains(value_type id) const -> bool
+    {
+        return std::find(cbegin(), cend(), id) != cend();
+    }
+
     auto ObjQueue::raw() const -> const ::Queue*
     {
         return &m_queue;
diff --git a/libmamba/src/solv-cpp/queue.hpp b/libmamba/src/solv-cpp/queue.hpp
index edc546016e..ddc8589601 100644
--- a/libmamba/src/solv-cpp/queue.hpp
+++ b/libmamba/src/solv-cpp/queue.hpp
@@ -84,6 +84,8 @@ namespace mamba::solv
         auto data() -> pointer;
         auto data() const -> const_pointer;
 
+        auto contains(value_type id) const -> bool;
+
         template <template <typename, typename...> class C>
         auto as() -> C<value_type>;
 
diff --git a/libmamba/tests/src/solv-cpp/test_queue.cpp b/libmamba/tests/src/solv-cpp/test_queue.cpp
index 0de9211f49..a5f2996f4a 100644
--- a/libmamba/tests/src/solv-cpp/test_queue.cpp
+++ b/libmamba/tests/src/solv-cpp/test_queue.cpp
@@ -205,4 +205,11 @@ TEST_SUITE("ObjQueue")
         q2.reserve(10);
         CHECK_EQ(q1, q2);
     }
+
+    TEST_CASE("contains")
+    {
+        const auto q = ObjQueue{ 1, 9, 3 };
+        CHECK(q.contains(3));
+        CHECK_FALSE(q.contains(0));
+    }
 }
diff --git a/libmamba/tests/src/util/test_flat_set.cpp b/libmamba/tests/src/util/test_flat_set.cpp
index 487e3fac78..ff4b9daef5 100644
--- a/libmamba/tests/src/util/test_flat_set.cpp
+++ b/libmamba/tests/src/util/test_flat_set.cpp
@@ -96,4 +96,106 @@ TEST_SUITE("flat_set")
         s.insert(6);
         CHECK_EQ(s.front(), 6);
     }
+
+    TEST_CASE("Set operations")
+    {
+        const auto s1 = flat_set<int>({ 1, 3, 4, 5 });
+        const auto s2 = flat_set<int>({ 3, 5 });
+        const auto s3 = flat_set<int>({ 4, 6 });
+
+        SUBCASE("Disjoint")
+        {
+            CHECK(s1.is_disjoint_of(flat_set<int>{}));
+            CHECK_FALSE(s1.is_disjoint_of(s1));
+            CHECK_FALSE(s1.is_disjoint_of(s2));
+            CHECK_FALSE(s1.is_disjoint_of(s3));
+            CHECK(s2.is_disjoint_of(s3));
+            CHECK(s3.is_disjoint_of(s2));
+        }
+
+        SUBCASE("Subset")
+        {
+            CHECK_LE(s1, s1);
+            CHECK_FALSE(s1 < s1);
+            CHECK_LE(flat_set<int>{}, s1);
+            CHECK_LT(flat_set<int>{}, s1);
+            CHECK_FALSE(s1 <= s2);
+            CHECK_FALSE(s1 <= flat_set<int>{});
+            CHECK_LE(flat_set<int>{ 1, 4 }, s1);
+            CHECK_LT(flat_set<int>{ 1, 4 }, s1);
+            CHECK_LE(s2, s1);
+            CHECK_LT(s2, s1);
+        }
+
+        SUBCASE("Superset")
+        {
+            CHECK_GE(s1, s1);
+            CHECK_FALSE(s1 > s1);
+            CHECK_GE(s1, flat_set<int>{});
+            CHECK_GT(s1, flat_set<int>{});
+            CHECK_FALSE(s2 >= s1);
+            CHECK_FALSE(flat_set<int>{} >= s1);
+            CHECK_GE(s1, flat_set<int>{ 1, 4 });
+            CHECK_GT(s1, flat_set<int>{ 1, 4 });
+            CHECK_GE(s1, s2);
+            CHECK_GT(s1, s2);
+        }
+
+        SUBCASE("Union")
+        {
+            CHECK_EQ(s1 | s1, s1);
+            CHECK_EQ(s1 | s2, s1);
+            CHECK_EQ(s2 | s1, s1 | s2);
+            CHECK_EQ(s1 | s3, flat_set<int>{ 1, 3, 4, 5, 6 });
+            CHECK_EQ(s3 | s1, s1 | s3);
+            CHECK_EQ(s2 | s3, flat_set<int>{ 3, 4, 5, 6 });
+            CHECK_EQ(s3 | s2, s2 | s3);
+        }
+
+        SUBCASE("Intersection")
+        {
+            CHECK_EQ(s1 & s1, s1);
+            CHECK_EQ(s1 & s2, s2);
+            CHECK_EQ(s2 & s1, s1 & s2);
+            CHECK_EQ(s1 & s3, flat_set<int>{ 4 });
+            CHECK_EQ(s3 & s1, s1 & s3);
+            CHECK_EQ(s2 & s3, flat_set<int>{});
+            CHECK_EQ(s3 & s2, s2 & s3);
+        }
+
+        SUBCASE("Difference")
+        {
+            CHECK_EQ(s1 - s1, flat_set<int>{});
+            CHECK_EQ(s1 - s2, flat_set<int>{ 1, 4 });
+            CHECK_EQ(s2 - s1, flat_set<int>{});
+            CHECK_EQ(s1 - s3, flat_set<int>{ 1, 3, 5 });
+            CHECK_EQ(s3 - s1, flat_set<int>{ 6 });
+            CHECK_EQ(s2 - s3, s2);
+            CHECK_EQ(s3 - s2, s3);
+        }
+
+        SUBCASE("Symetric difference")
+        {
+            CHECK_EQ(s1 ^ s1, flat_set<int>{});
+            CHECK_EQ(s1 ^ s2, flat_set<int>{ 1, 4 });
+            CHECK_EQ(s2 ^ s1, s1 ^ s2);
+            CHECK_EQ(s1 ^ s3, flat_set<int>{ 1, 3, 5, 6 });
+            CHECK_EQ(s3 ^ s1, s1 ^ s3);
+            CHECK_EQ(s2 ^ s3, flat_set<int>{ 3, 4, 5, 6 });
+            CHECK_EQ(s3 ^ s2, s2 ^ s3);
+        }
+
+        SUBCASE("Algebra")
+        {
+            for (const auto& u : { s1, s2, s3 })
+            {
+                for (const auto& v : { s1, s2, s3 })
+                {
+                    CHECK_EQ((u - v) | (v - u) | (u & v), u | v);
+                    CHECK_EQ((u ^ v) | (u & v), u | v);
+                    CHECK_EQ((u | v) - (u & v), u ^ v);
+                }
+            }
+        }
+    }
 }