From 6c9b43787c50f5208f11eb5883882129d362cbca Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Tue, 28 May 2024 19:38:03 +0000
Subject: [PATCH] [Bugfix][Support] Fix copy constructor for
 support::OrderedSet

Prior to this commit, the `support::OrderedSet<T>` utility used the
default copy constructor and copy assignment, which would copy both
the `OrderedSet::elements_` and `OrderedSet::elem_to_iter_` members.
While this is the correct behavior for `elements_`, the copy of
`elem_to_iter_` would contain references to the original's `element_`,
rather than to its own.

While `elem_to_iter_` is used in both `OrderedSet::push_back` and
`OrderedSet::erase`, the implementation of `OrderedSet::push_back`
only depends on the keys used in `elem_to_iter_`, and does not depend
on the values stored.  As a result, this bug could go undetected for
append-only usage, which is the most frequent use of `OrderedSet`.

This commit updates `support::OrderedSet` to have an explicit copy
constructor and copy assignment.  Only the `std::list<T> elements_`
member may be copied, while the `elem_to_iter_` must instead be
rebuilt.
---
 src/support/ordered_set.h | 31 +++++++++++++++++++++++++++----
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/support/ordered_set.h b/src/support/ordered_set.h
index 741f0b18e6b9..11acb8c3fef5 100644
--- a/src/support/ordered_set.h
+++ b/src/support/ordered_set.h
@@ -54,11 +54,28 @@ class OrderedSet {
  public:
   OrderedSet() = default;
 
+  /* \brief Explicit copy constructor
+   *
+   * The default copy constructor would copy both `elements_` and
+   * `elem_to_iter_`.  While this is the correct behavior for
+   * `elements_`, the copy of `elem_to_iter_` would contain references
+   * to the original's `element_`, rather than to its own
+   */
+  OrderedSet(const OrderedSet<T>& other) : elements_(other.elements_) { InitElementToIter(); }
+
+  /* \brief Explicit copy assignment
+   *
+   * Implemented in terms of the copy constructor, and the default
+   * move assignment.
+   */
+  OrderedSet& operator=(const OrderedSet<T>& other) { return *this = OrderedSet(other); }
+
+  OrderedSet(OrderedSet<T>&&) = default;
+  OrderedSet& operator=(OrderedSet<T>&&) = default;
+
   template <typename Iter>
-  OrderedSet(Iter begin, Iter end) {
-    for (auto it = begin; it != end; it++) {
-      push_back(*it);
-    }
+  OrderedSet(Iter begin, Iter end) : elements_(begin, end) {
+    InitElementToIter();
   }
 
   void push_back(const T& t) {
@@ -90,6 +107,12 @@ class OrderedSet {
   auto empty() const { return elements_.empty(); }
 
  private:
+  void InitElementToIter() {
+    for (auto it = elements_.begin(); it != elements_.end(); it++) {
+      elem_to_iter_[*it] = it;
+    }
+  }
+
   std::list<T> elements_;
   typename detail::OrderedSetLookupType<T>::MapType elem_to_iter_;
 };