Utils: Refactor and Test Updates

lib/utils/include/utils/containers.decl.h

@@ -30,27 +30,16 @@ bool contains_l(bidict<K, V> const &m, K const &k);
 template <typename K, typename V>
 bool contains_r(bidict<K, V> const &m, V const &v);
 
-template <typename K, typename V, typename F>
-std::unordered_map<K, V> filter_values(std::unordered_map<K, V> const &m,
-                                       F const &f);
-
 template <typename Container, typename Element>
 std::optional<std::size_t> index_of(Container const &c, Element const &e);
 
-template <typename K, typename V>
-std::unordered_map<K, V> restrict_keys(std::unordered_map<K, V> const &m,
-                                       std::unordered_set<K> const &mask);
-
 template <typename K, typename V>
 std::unordered_map<K, V> merge_maps(std::unordered_map<K, V> const &lhs,
                                     std::unordered_map<K, V> const &rhs);
 
 template <typename K, typename V>
 bidict<K, V> merge_maps(bidict<K, V> const &lhs, bidict<K, V> const &rhs);
 
-template <typename E>
-std::optional<E> at_idx(std::vector<E> const &v, size_t idx);
-
 template <typename K, typename V>
 std::function<V(K const &)> lookup_in(std::unordered_map<K, V> const &m);
 
@@ -61,32 +50,30 @@ template <typename L, typename R>
 std::function<L(R const &)> lookup_in_r(bidict<L, R> const &m);
 
 template <typename T>
-bool is_supserseteq_of(std::unordered_set<T> const &l,
-                       std::unordered_set<T> const &r);
-
-template <typename S, typename D>
-std::unordered_set<D>
-    map_over_unordered_set(std::function<D(S const &)> const &f,
-                           std::unordered_set<S> const &input);
-
-template <typename C>
-std::optional<typename C::value_type> maybe_get_only(C const &c);
+bool is_superseteq_of(std::unordered_set<T> const &l,
+                      std::unordered_set<T> const &r);
 
 template <typename Container, typename Function>
 std::optional<bool> optional_all_of(Container const &, Function const &);
 
 template <typename C>
 bool are_all_same(C const &c);
 
+template <typename In, typename F, typename Out>
+std::vector<Out> flatmap(std::vector<In> const &v, F const &f);
+
+template <typename In, typename F, typename Out>
+std::unordered_set<Out> flatmap(std::unordered_set<In> const &v, F const &f);
+template <typename Out, typename In>
+std::unordered_set<Out> flatmap_v2(std::unordered_set<In> const &v,
+                                   std::unordered_set<Out> (*f)(In const &));
+
 template <typename T, typename F>
 std::function<bool(T const &, T const &)> compare_by(F const &f);
 
 template <typename C>
 typename C::value_type maximum(C const &v);
 
-template <typename T>
-T reversed(T const &t);
-
 template <typename T>
 std::vector<T> value_all(std::vector<std::optional<T>> const &v);
 

lib/utils/include/utils/containers.h

@@ -5,17 +5,21 @@
 #include "required_core.h"
 #include "type_traits_core.h"
 #include "utils/bidict/bidict.h"
+#include "utils/containers/are_disjoint.h"
 #include "utils/containers/contains.h"
 #include "utils/containers/extend.h"
 #include "utils/containers/extend_vector.h"
 #include "utils/containers/filter.h"
 #include "utils/containers/intersection.h"
 #include "utils/containers/is_subseteq_of.h"
+#include "utils/containers/keys.h"
+#include "utils/containers/restrict_keys.h"
 #include "utils/containers/sorted.h"
 #include "utils/containers/transform.h"
 #include "utils/containers/vector_transform.h"
 #include "utils/exception.h"
 #include "utils/hash/pair.h"
+#include "utils/optional.h"
 #include "utils/type_traits.h"
 #include <algorithm>
 #include <cassert>
@@ -135,21 +139,11 @@ std::function<L(R const &)> lookup_in_r(bidict<L, R> const &m) {
 }
 
 template <typename T>
-bool is_supserseteq_of(std::unordered_set<T> const &l,
-                       std::unordered_set<T> const &r) {
+bool is_superseteq_of(std::unordered_set<T> const &l,
+                      std::unordered_set<T> const &r) {
   return is_subseteq_of<T>(r, l);
 }
 
-template <typename S, typename D>
-std::unordered_set<D>
-    map_over_unordered_set(std::function<D(S const &)> const &f,
-                           std::unordered_set<S> const &input) {
-  std::unordered_set<D> result;
-  std::transform(
-      input.cbegin(), input.cend(), std::inserter(result, result.begin()), f);
-  return result;
-}
-
 template <typename Container, typename Function>
 std::optional<bool> optional_all_of(Container const &container,
                                     Function const &func) {

lib/utils/include/utils/containers/vector_split.h

@@ -1,6 +1,7 @@
 #ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_CONTAINERS_VECTOR_SPLIT_H
 #define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_CONTAINERS_VECTOR_SPLIT_H
 
+#include <cassert>
 #include <vector>
 
 namespace FlexFlow {

lib/utils/include/utils/graph/dataflow_graph/dataflow_graph.h

@@ -43,8 +43,6 @@ struct DataflowGraph : virtual public DataflowGraphView {
 private:
   IDataflowGraph &get_interface();
   IDataflowGraph const &get_interface() const;
-
-  friend struct GraphInternal;
 };
 
 } // namespace FlexFlow

lib/utils/include/utils/graph/digraph/digraph.h

@@ -40,8 +40,6 @@ struct DiGraph : virtual DiGraphView {
 private:
   IDiGraph &get_ptr();
   IDiGraph const &get_ptr() const;
-
-  friend struct GraphInternal;
 };
 CHECK_WELL_BEHAVED_VALUE_TYPE_NO_EQ(DiGraph);
 

lib/utils/include/utils/graph/digraph/digraph_view.h

@@ -31,8 +31,6 @@ struct DiGraphView : virtual public GraphView {
 
 private:
   IDiGraphView const &get_ptr() const;
-
-  friend struct GraphInternal;
 };
 CHECK_WELL_BEHAVED_VALUE_TYPE_NO_EQ(DiGraphView);
 

lib/utils/include/utils/graph/multidigraph/multidigraph.h

@@ -40,8 +40,6 @@ struct MultiDiGraph : virtual public MultiDiGraphView {
 private:
   IMultiDiGraph &get_interface();
   IMultiDiGraph const &get_interface() const;
-
-  friend struct GraphInternal;
 };
 
 } // namespace FlexFlow

lib/utils/include/utils/graph/node/graph.h

@@ -31,8 +31,6 @@ struct Graph : virtual GraphView {
 private:
   IGraph const &get_ptr() const;
   IGraph &get_ptr();
-
-  friend struct GraphInternal;
 };
 
 } // namespace FlexFlow

lib/utils/include/utils/graph/node/graph_view.h

@@ -22,8 +22,6 @@ struct GraphView {
   GraphView();
   cow_ptr_t<IGraphView> ptr;
   GraphView(cow_ptr_t<IGraphView> ptr);
-
-  friend struct GraphInternal;
 };
 
 } // namespace FlexFlow

lib/utils/include/utils/graph/undirected/undirected_edge.h

@@ -31,4 +31,9 @@ struct hash<::FlexFlow::UndirectedEdge> {
 
 } // namespace std
 
+namespace FlexFlow {
+std::string format_as(UndirectedEdge const &);
+std::ostream &operator<<(std::ostream &, UndirectedEdge const &);
+} // namespace FlexFlow
+
 #endif

lib/utils/include/utils/graph/undirected/undirected_graph.h

@@ -34,8 +34,6 @@ struct UndirectedGraph : virtual UndirectedGraphView {
 
   using UndirectedGraphView::UndirectedGraphView;
 
-  friend struct GraphInternal;
-
 private:
   IUndirectedGraph const &get_ptr() const;
   IUndirectedGraph &get_ptr();

lib/utils/include/utils/graph/undirected/undirected_graph_view.h

@@ -29,8 +29,6 @@ struct UndirectedGraphView : virtual GraphView {
 
   using GraphView::GraphView;
 
-  friend struct GraphInternal;
-
 private:
   IUndirectedGraphView const &get_ptr() const;
 };

lib/utils/src/utils/graph/undirected/undirected_edge.cc

@@ -1,5 +1,6 @@
 #include "utils/graph/undirected/undirected_edge.h"
 #include "utils/hash/tuple.h"
+#include <sstream>
 
 namespace FlexFlow {
 
@@ -38,3 +39,17 @@
 }
 
 } // namespace std
+
+namespace FlexFlow {
+std::string format_as(UndirectedEdge const &x) {
+  std::ostringstream oss;
+  oss << "<UndirectedEdge";
+  oss << " smaller=" << x.smaller;
+  oss << " bigger=" << x.bigger;
+  oss << ">";
+  return oss.str();
+}
+std::ostream &operator<<(std::ostream &s, UndirectedEdge const &x) {
+  return s << fmt::to_string(x);
+}
+} // namespace FlexFlow

lib/utils/test/src/test_algorithms.cc

@@ -75,16 +75,6 @@ TEST_SUITE(FF_TEST_SUITE) {
       CHECK(result == expected_result);
     }
 
-    SUBCASE("get_dominators") {
-      std::unordered_map<Node, std::unordered_set<Node>> expected = {
-          {n[0], {n[0]}},
-          {n[1], {n[0], n[1]}},
-          {n[2], {n[0], n[2]}},
-          {n[3], {n[0], n[3]}},
-      };
-      CHECK(get_dominators(g) == expected);
-    }
-
     SUBCASE("get_sinks") {
       auto expected = std::unordered_set<Node>{n[2], n[3]};
       CHECK(get_sinks(g) == expected);