From cd7fef3efb51bd9f61a278d08fe3dbf805f6b08b Mon Sep 17 00:00:00 2001
From: lambda7xx <lambda7xx@gmail.com>
Date: Mon, 17 Jul 2023 07:50:58 +0000
Subject: [PATCH] implement the query_values and query_keys for
 std::unordered_map and bidict fix the bug of containers

---
 lib/utils/include/utils/containers.h          | 34 +++++++++++++++++--
 lib/utils/include/utils/graph/query_set.h     | 30 ++++++++++++++--
 lib/utils/src/graph/adjacency_multidigraph.cc |  3 +-
 lib/utils/src/graph/views.cc                  |  2 +-
 4 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/lib/utils/include/utils/containers.h b/lib/utils/include/utils/containers.h
index c75876ecf1..ae274f473a 100644
--- a/lib/utils/include/utils/containers.h
+++ b/lib/utils/include/utils/containers.h
@@ -110,6 +110,11 @@ typename It::value_type product(It begin, It end) {
       });
 }
 
+// template <typename Container>
+// bool contains(Container const &c, typename Container::C::key_type const &e) {
+//   return find<Container>(c, e) != c.cend();
+// }
+
 template <typename Container>
 bool contains(Container const &c, typename Container::value_type const &e) {
   return find<Container>(c, e) != c.cend();
@@ -149,7 +154,7 @@ template <typename K,
           typename K2 = decltype(std::declval<F>()(std::declval<K>()))>
 bidict<K2, V> map_keys(bidict<K, V> const &m, F const &f) {
   bidict<K2, V> result;
-  for (auto const &kv : f) {
+  for (auto const &kv : m) {
     result.equate(f(kv.first), kv.second);
   }
   return result;
@@ -159,7 +164,17 @@ template <typename K, typename V, typename F>
 std::unordered_map<K, V> filter_keys(std::unordered_map<K, V> const &m,
                                      F const &f) {
   std::unordered_map<K, V> result;
-  for (auto const &kv : f) {
+  for (auto const &kv : m) {
+    if (f(kv.first)) {
+      result.insert(kv);
+    }
+  }
+  return result;
+}
+
+template <typename K, typename V, typename F> std::unordered_map<K, V> filter_keys(bidict<K, V> const & m, F const & f) {
+  std::unordered_map<K, V> result;
+  for (auto const &kv : m) {
     if (f(kv.first)) {
       result.insert(kv);
     }
@@ -167,6 +182,18 @@ std::unordered_map<K, V> filter_keys(std::unordered_map<K, V> const &m,
   return result;
 }
 
+
+template <typename K, typename V, typename F> std::unordered_map<K, V> filter_values(bidict<K, V> const & m, F const & f) {
+  std::unordered_map<K, V> result;
+  for (auto const &kv : m) {
+    if (f(kv.second)) {
+      result.insert(kv);
+    }
+  }
+  return result;
+}
+
+
 template <typename K,
           typename V,
           typename F,
@@ -204,6 +231,9 @@ std::unordered_map<K, V> filter_values(std::unordered_map<K, V> const &m,
   return result;
 }
 
+
+
+
 template <typename C>
 std::vector<typename C::key_type> keys(C const &c) {
   std::vector<typename C::key_type> result;
diff --git a/lib/utils/include/utils/graph/query_set.h b/lib/utils/include/utils/graph/query_set.h
index 957f73cb1a..d16099a080 100644
--- a/lib/utils/include/utils/graph/query_set.h
+++ b/lib/utils/include/utils/graph/query_set.h
@@ -4,6 +4,7 @@
 #include "utils/containers.h"
 #include "utils/exception.h"
 #include "utils/optional.h"
+#include "utils/bidict.h"
 #include <unordered_set>
 
 namespace FlexFlow {
@@ -71,15 +72,40 @@ template <typename C,
           typename V = typename C::mapped_type>
 std::unordered_map<K, V> query_keys(query_set<K> const &q, C const &m) {
   std::cout<<"3"<<std::endl;
-  NOT_IMPLEMENTED();
+  std::unordered_set<K> q_set = allowed_values(q);
+
+  auto filter_lambda = [&q_set](const K& key) {
+    return q_set.find(key) != q_set.end();
+  };
+
+  return filter_keys(m, filter_lambda);
 }//TODO
 
+template <typename K,
+          typename V>
+std::unordered_map<K,V> query_keys(query_set<V> const &q, bidict<K, V> const &m) {
+  std::unordered_set<V> q_set = allowed_values(q);
+
+  auto filter_lambda = [&q_set](const V& value) {
+    return q_set.find(value) != q_set.end();
+  };
+
+  return  filter_values(m, filter_lambda);
+}
+
+
 template <typename C,
           typename K = typename C::key_type,
           typename V = typename C::mapped_type>
 std::unordered_map<K, V> query_values(query_set<V> const &q, C const &m) {
   std::cout<<"4"<<std::endl;
-  NOT_IMPLEMENTED();
+  std::unordered_set<V> q_set = allowed_values(q);
+
+  auto filter_lambda = [&q_set](const V& value) {
+    return q_set.find(value) != q_set.end();
+  };
+
+  return filter_values(m, filter_lambda);
 }
 
 template <typename T>
diff --git a/lib/utils/src/graph/adjacency_multidigraph.cc b/lib/utils/src/graph/adjacency_multidigraph.cc
index ce9fd33d8d..0b8ee7f4f1 100644
--- a/lib/utils/src/graph/adjacency_multidigraph.cc
+++ b/lib/utils/src/graph/adjacency_multidigraph.cc
@@ -42,10 +42,11 @@ void AdjacencyMultiDiGraph::remove_edge(MultiDiEdge const &e) {
 std::unordered_set<MultiDiEdge>
     AdjacencyMultiDiGraph::query_edges(MultiDiEdgeQuery const &q) const {
   std::unordered_set<MultiDiEdge> result;
-  for (auto const &src_kv : query_keys(q.srcs, this->adjacency)) {
+  for (auto const &src_kv : query_keys(q.srcs, this->adjacency)) { //query_keys(q.srcs, this->adjacency) return map<node, map<node, map<nodeport, set<nodeport>>>>
     for (auto const &dst_kv : query_keys(q.dsts, src_kv.second)) {
       for (auto const &srcIdx_kv : query_keys(q.srcIdxs, dst_kv.second)) {
         for (auto const &dstIdx : apply_query(q.dstIdxs, srcIdx_kv.second)) {
+          std::cout<<"add edge "<<src_kv.first.value()<<" "<<dst_kv.first.value()<<" "<<srcIdx_kv.first.value()<<" "<<dstIdx.value()<<std::endl;
           result.insert({src_kv.first, dst_kv.first, srcIdx_kv.first, dstIdx});
         }
       }
diff --git a/lib/utils/src/graph/views.cc b/lib/utils/src/graph/views.cc
index 98f2ab7271..ca52d2afdc 100644
--- a/lib/utils/src/graph/views.cc
+++ b/lib/utils/src/graph/views.cc
@@ -147,7 +147,7 @@ JoinedNodeView::JoinedNodeView(GraphView const &lhs, GraphView const &rhs) {
 
 std::unordered_set<Node>
     JoinedNodeView::query_nodes(NodeQuery const &query) const {
-  return unique(values(query_keys(query.nodes, this->mapping)));
+  return unique(values(query_keys(query.nodes, this->mapping)));//query_keys(query.nodes, this->mapping)->std::unordered_map<x, Node>
 }
 
 std::pair<std::unordered_set<Node>, std::unordered_set<Node>>