Improve bidict implementation and tests

lib/utils/include/utils/bidict/algorithms/bidict_filter_keys.h

@@ -10,7 +10,7 @@ bidict<K, V> bidict_filter_keys(bidict<K, V> const &m, F &&f) {
   bidict<K, V> result;
   for (auto const &kv : m) {
     if (f(kv.first)) {
-      result.equate(kv);
+      result.equate_strict(kv);
     }
   }
   return result;

lib/utils/include/utils/bidict/algorithms/bidict_filter_values.h

@@ -10,7 +10,7 @@ bidict<K, V> bidict_filter_values(bidict<K, V> const &m, F &&f) {
   bidict<K, V> result;
   for (auto const &kv : m) {
     if (f(kv.second)) {
-      result.equate(kv);
+      result.equate_strict(kv);
     }
   }
   return result;

lib/utils/include/utils/bidict/algorithms/bidict_filtrans_keys.h

@@ -14,7 +14,7 @@ bidict<K2, V> bidict_filtrans_keys(bidict<K, V> const &m, F &&f) {
   for (auto const &[k, v] : m) {
     std::optional<K2> new_k = f(k);
     if (new_k.has_value()) {
-      result.equate(new_k.value(), v);
+      result.equate_strict(new_k.value(), v);
     }
   }
   return result;

lib/utils/include/utils/bidict/algorithms/bidict_filtrans_values.h

@@ -14,7 +14,7 @@ bidict<K, V2> bidict_filtrans_values(bidict<K, V> const &m, F &&f) {
   for (auto const &[k, v] : m) {
     std::optional<V2> new_v = f(v);
     if (new_v.has_value()) {
-      result.equate(k, new_v.value());
+      result.equate_strict(k, new_v.value());
     }
   }
   return result;

lib/utils/include/utils/bidict/algorithms/bidict_from_enumerating.h

@@ -29,7 +29,7 @@ bidict<nonnegative_int, T>
   bidict<nonnegative_int, T> result;
   nonnegative_int idx = 0_n;
   for (T const &t : s) {
-    result.equate(idx, t);
+    result.equate_strict(idx, t);
     idx++;
   }
 

lib/utils/include/utils/bidict/algorithms/bidict_from_map.h

@@ -7,22 +7,12 @@ namespace FlexFlow {
 
 template <typename L, typename R>
 bidict<L, R> bidict_from_map(std::map<L, R> const &m) {
-  bidict<L, R> result;
-  for (auto const &[k, v] : m) {
-    ASSERT(!result.contains_r(v));
-    result.equate({k, v});
-  }
-  return result;
+  return bidict<L, R>{m.begin(), m.end()};
 }
 
 template <typename L, typename R>
 bidict<L, R> bidict_from_map(std::unordered_map<L, R> const &m) {
-  bidict<L, R> result;
-  for (auto const &[k, v] : m) {
-    ASSERT(!result.contains_r(v));
-    result.equate({k, v});
-  }
-  return result;
+  return bidict<L, R>{m.begin(), m.end()};
 }
 
 } // namespace FlexFlow

lib/utils/include/utils/bidict/algorithms/bidict_from_unstructured_relation.h

@@ -44,11 +44,7 @@ bidict<L, R> bidict_from_unstructured_relation(
     ASSERT(duplicated_element_counts.empty(), duplicated_element_counts);
   }
 
-  bidict<L, R> result;
-  for (auto const &lr : relation) {
-    result.equate_strict(lr);
-  }
-  return result;
+  return bidict<L, R>{relation.begin(), relation.end()};
 }
 
 } // namespace FlexFlow

lib/utils/include/utils/bidict/algorithms/exhaustive_relational_join.h

@@ -16,7 +16,7 @@ bidict<T1, T3> exhaustive_relational_join(bidict<T1, T2> const &fst,
   bidict<T1, T3> result;
 
   for (auto const &[v1, v2] : fst) {
-    result.equate({v1, snd.at_l(v2)});
+    result.equate_strict({v1, snd.at_l(v2)});
   }
 
   return result;

lib/utils/include/utils/bidict/algorithms/transform.h

@@ -13,7 +13,7 @@ template <typename K,
 bidict<K2, V2> transform(bidict<K, V> const &m, F &&f) {
   bidict<K2, V2> result;
   for (auto const &[k, v] : m) {
-    result.equate(f(k, v));
+    result.equate_strict(f(k, v));
   }
   return result;
 }

lib/utils/include/utils/bidict/bidict.h

@@ -30,7 +30,7 @@ struct bidict {
   template <typename InputIt>
   bidict(InputIt first, InputIt last) {
     for (auto it = first; it != last; it++) {
-      this->equate(it->first, it->second);
+      this->equate_strict(it->first, it->second);
     }
   }
 
@@ -47,37 +47,40 @@ struct bidict {
   }
 
   void erase_l(L const &l) {
-    this->fwd_map.erase(l);
-    for (auto const &kv : this->bwd_map) {
-      if (kv.second == l) {
-        bwd_map.erase(kv.first);
-        break;
-      }
+    if (this->contains_l(l)) {
+      R r = this->at_l(l);
+      this->fwd_map.erase(l);
+      this->bwd_map.erase(r);
     }
   }
 
   void erase_r(R const &r) {
-    this->bwd_map.erase(r);
-    for (auto const &kv : this->fwd_map) {
-      if (kv.second == r) {
-        fwd_map.erase(kv.first);
-        break;
-      }
+    if (this->contains_r(r)) {
+      L l = this->at_r(r);
+      this->fwd_map.erase(l);
+      this->bwd_map.erase(r);
     }
   }
 
   void equate(L const &l, R const &r) {
-    fwd_map.insert({l, r});
-    bwd_map.insert({r, l});
+    bool contains_l = this->contains_l(l);
+    bool contains_r = this->contains_r(r);
 
-    this->check_invariants();
+    if (contains_l != contains_r || (contains_l && this->at_l(l) != r)) {
+      this->erase_l(l);
+      this->erase_r(r);
+      contains_l = contains_r = false;
+    }
+
+    if (!contains_l) {
+      ASSERT(!contains_r);
+      fwd_map.insert({l, r});
+      bwd_map.insert({r, l});
+    }
   }
 
   void equate(std::pair<L, R> const &lr) {
-    fwd_map.insert(lr);
-    bwd_map.insert({lr.second, lr.first});
-
-    this->check_invariants();
+    this->equate(lr.first, lr.second);
   }
 
   void equate_strict(L const &l, R const &r) {
@@ -86,7 +89,8 @@ struct bidict {
     if (this->contains_l(l)) {
       ASSERT(this->at_l(l) == r);
     } else {
-      this->equate(l, r);
+      fwd_map.insert({l, r});
+      bwd_map.insert({r, l});
     }
   }
 
@@ -238,7 +242,9 @@ struct bidict {
   }
 
   bidict(std::map<L, R> const &fwd_map, std::map<R, L> const &bwd_map)
-      : fwd_map(fwd_map), bwd_map(bwd_map) {}
+      : fwd_map(fwd_map), bwd_map(bwd_map) {
+    this->check_invariants();
+  }
 
   bool operator<(bidict<L, R> const &other) const {
     return this->fwd_map < other.fwd_map;

lib/utils/test/src/utils/bidict/bidict.cc

@@ -29,26 +29,92 @@ TEST_SUITE(FF_TEST_SUITE) {
     }
 
     SUBCASE("L type is not the same as R type") {
-      bidict<int, std::string> dict;
-      dict.equate(1, "one");
-      dict.equate(2, "two");
+      bidict<int, std::string> bd;
+      bd.equate(1, "one");
+      bd.equate(2, "two");
 
       SUBCASE("bidict::contains_l") {
-        CHECK(dict.contains_l(1));
-        CHECK_FALSE(dict.contains_l(3));
+        CHECK(bd.contains_l(1));
+        CHECK_FALSE(bd.contains_l(3));
       }
 
       SUBCASE("bidict::contains_r") {
-        CHECK(dict.contains_r("one"));
-        CHECK_FALSE(dict.contains_r("three"));
+        CHECK(bd.contains_r("one"));
+        CHECK_FALSE(bd.contains_r("three"));
+      }
+    }
+
+    SUBCASE("bidict::bidict(std::initializer_list<std::pair<L, R>>)") {
+      bidict<int, std::string> bd{{1, "one"}, {2, "two"}};
+      CHECK(bd.contains_l(1));
+      CHECK_FALSE(bd.contains_l(3));
+
+      SUBCASE("invalid mapping") {
+        CHECK_THROWS(bidict<int, std::string>{{1, "one"}, {2, "one"}});
       }
     }
 
+    SUBCASE("bidict::bidict(InputIt)") {
+      std::vector<std::pair<int, std::string>> pairs = {{1, "one"}, {2, "two"}};
+      bidict<int, std::string> bd{pairs.begin(), pairs.end()};
+      CHECK(bd.contains_l(1));
+      CHECK_FALSE(bd.contains_l(3));
+
+      SUBCASE("invalid mapping") {
+        std::vector<std::pair<int, std::string>> bad_pairs = {{1, "one"},
+                                                              {2, "one"}};
+        CHECK_THROWS(
+            bidict<int, std::string>{bad_pairs.begin(), bad_pairs.end()});
+      }
+    }
+
+    SUBCASE("bidict::bidict(std::map<L, R> const &, std::map<R, L> const &)") {
+      std::map<int, std::string> fwd = {{1, "one"}, {2, "two"}};
+      std::map<std::string, int> bwd = {{"one", 1}, {"two", 2}};
+      bidict<int, std::string> bd{fwd, bwd};
+      CHECK(bd.contains_l(1));
+      CHECK_FALSE(bd.contains_l(3));
+
+      SUBCASE("invalid mapping") {
+        std::map<int, std::string> bad_fwd = {{1, "one"}, {2, "one"}};
+        std::map<std::string, int> bad_bwd = {{"one", 1}};
+        CHECK_THROWS(bidict<int, std::string>{bad_fwd, bad_bwd});
+      }
+    }
+
+    SUBCASE("bidict::erase_l") {
+      dict.erase_l(1);
+      CHECK(dict.size() == 1);
+      CHECK_THROWS(dict.at_l(1));
+      CHECK(dict.at_r("two") == 2);
+    }
+
+    SUBCASE("bidict::erase_r") {
+      dict.erase_r("one");
+      CHECK(dict.size() == 1);
+      CHECK_THROWS(dict.at_r("one"));
+      CHECK(dict.at_l(2) == "two");
+    }
+
     SUBCASE("bidict::equate") {
       CHECK(dict.at_l(1) == "one");
       CHECK(dict.at_r("one") == 1);
       CHECK(dict.at_l(2) == "two");
       CHECK(dict.at_r("two") == 2);
+
+      dict.equate(1, "three");
+      CHECK(dict.at_l(1) == "three");
+      CHECK(dict.at_r("three") == 1);
+      CHECK_THROWS(dict.at_r("one"));
+      CHECK(dict.at_l(2) == "two");
+      CHECK(dict.at_r("two") == 2);
+
+      dict.equate(3, "three");
+      CHECK(dict.at_l(3) == "three");
+      CHECK(dict.at_r("three") == 3);
+      CHECK_THROWS(dict.at_l(1));
+      CHECK(dict.at_l(2) == "two");
+      CHECK(dict.at_r("two") == 2);
     }
 
     SUBCASE("bidict::equate_strict") {
@@ -60,34 +126,46 @@ TEST_SUITE(FF_TEST_SUITE) {
       CHECK(dict.at_r("three") == 3);
     }
 
-    SUBCASE("bidict::erase_l") {
-      dict.erase_l(1);
-      CHECK(dict.size() == 1);
-      CHECK_THROWS(dict.at_l(1));
-      CHECK(dict.at_r("two") == 2);
+    SUBCASE("bidict::operator==") {
+      bidict<int, std::string> bd{{1, "one"}, {2, "two"}};
+      bidict<int, std::string> bd2{{1, "one"}, {3, "three"}};
+      CHECK(dict == bd);
+      CHECK_FALSE(dict == bd2);
     }
 
-    SUBCASE("bidict::erase_r") {
-      dict.erase_r("one");
-      CHECK(dict.size() == 1);
-      CHECK_THROWS(dict.at_r("one"));
-      CHECK(dict.at_l(2) == "two");
+    SUBCASE("bidict::operator!=") {
+      bidict<int, std::string> bd{{1, "one"}, {2, "two"}};
+      bidict<int, std::string> bd2{{1, "one"}, {3, "three"}};
+      CHECK_FALSE(dict != bd);
+      CHECK(dict != bd2);
     }
 
-    SUBCASE("bidict::reversed") {
-      bidict<std::string, int> reversed_dict = dict.reversed();
-      CHECK(reversed_dict.at_l("one") == 1);
-      CHECK(reversed_dict.at_r(2) == "two");
+    SUBCASE("bidict::at_l") {
+      CHECK(dict.at_l(1) == "one");
+      CHECK_THROWS(dict.at_l(3));
+    }
+
+    SUBCASE("bidict::at_r") {
+      CHECK(dict.at_r("one") == 1);
+      CHECK_THROWS(dict.at_r("three"));
+    }
+
+    SUBCASE("bidict::left_values") {
+      CHECK(dict.left_values() == std::set<int>{1, 2});
+    }
+
+    SUBCASE("bidict::right_values") {
+      CHECK(dict.right_values() == std::set<std::string>{"one", "two"});
     }
 
     SUBCASE("bidict::size") {
       CHECK(dict.size() == 2);
     }
 
-    SUBCASE("implicitly convert to std::map") {
-      std::map<int, std::string> res = dict;
-      std::map<int, std::string> expected = {{1, "one"}, {2, "two"}};
-      CHECK(res == expected);
+    SUBCASE("bidict::empty") {
+      CHECK_FALSE(dict.empty());
+      bidict<int, std::string> empty{};
+      CHECK(empty.empty());
     }
 
     SUBCASE("bidict::begin") {
@@ -102,6 +180,18 @@ TEST_SUITE(FF_TEST_SUITE) {
       CHECK_WITHOUT_STRINGIFY(it == dict.end());
     }
 
+    SUBCASE("bidict::reversed") {
+      bidict<std::string, int> reversed_dict = dict.reversed();
+      CHECK(reversed_dict.at_l("one") == 1);
+      CHECK(reversed_dict.at_r(2) == "two");
+    }
+
+    SUBCASE("implicitly convert to std::map") {
+      std::map<int, std::string> res = dict;
+      std::map<int, std::string> expected = {{1, "one"}, {2, "two"}};
+      CHECK(res == expected);
+    }
+
     SUBCASE("fmt::to_string(bidict<int, std::string>)") {
       std::string result = fmt::to_string(dict);
       std::string correct = fmt::to_string(dict.as_map());