From 919febe4cc731dc318ed6e3a01787e6d764a7d51 Mon Sep 17 00:00:00 2001
From: Nathan Hughes <na26933@mit.edu>
Date: Fri, 19 Jul 2024 13:11:14 +0000
Subject: [PATCH] (wip) start experimenting with moving merge to edge container

---
 include/spark_dsg/edge_container.h |  3 ++
 src/edge_container.cpp             | 61 +++++++++++++++++++++++++++++-
 src/scene_graph_layer.cpp          | 18 +--------
 3 files changed, 64 insertions(+), 18 deletions(-)

diff --git a/include/spark_dsg/edge_container.h b/include/spark_dsg/edge_container.h
index 9519521..7703317 100644
--- a/include/spark_dsg/edge_container.h
+++ b/include/spark_dsg/edge_container.h
@@ -132,6 +132,9 @@ struct EdgeContainer {
 
   void setStale();
 
+  void updateFrom(const EdgeContainer& other,
+                  const std::map<NodeId, NodeId>* merges = nullptr);
+
   Edges edges;
   EdgeStatusMap edge_status;
   mutable std::map<EdgeKey, bool> stale_edges;
diff --git a/src/edge_container.cpp b/src/edge_container.cpp
index dcb2c39..a391a4b 100644
--- a/src/edge_container.cpp
+++ b/src/edge_container.cpp
@@ -142,6 +142,66 @@ void EdgeContainer::setStale() {
   }
 }
 
+NodeId getMergedId(NodeId node, const std::map<NodeId, NodeId>* merges) {
+  if (!merges) {
+    return node;
+  }
+
+  auto iter = merges->find(node);
+  return iter == merges->end() ? node : iter->second;
+}
+
+void EdgeContainer::updateFrom(const EdgeContainer& other,
+                               const std::map<NodeId, NodeId>* merges) {
+  for (const auto& id_edge_pair : other_layer.edges_.edges) {
+    const auto& edge = id_edge_pair.second;
+    NodeId new_source = config.getMergedId(edge.source);
+    NodeId new_target = config.getMergedId(edge.target);
+    if (new_source == new_target) {
+      continue;
+    }
+
+    // NOTE(nathan) not really necessary but saves the clone call
+    if (hasEdge(new_source, new_target)) {
+      // TODO(nathan) clone attributes
+      continue;
+    }
+
+    insertEdge(new_source, new_target, edge.info->clone());
+  }
+
+  std::vector<EdgeKey> removed_edges;
+  other_layer->edges_.getRemoved(removed_edges, config.clear_removed);
+  for (const auto& removed_edge : removed_edges) {
+    NodeId new_source = config.getMergedId(removed_edge.k1);
+    NodeId new_target = config.getMergedId(removed_edge.k2);
+    if (new_source == new_target) {
+      continue;
+    }
+
+    if (other_layer->hasEdge(new_source, new_target)) {
+      continue;  // edge still exists through merged node
+    }
+
+    layers_[l_id]->removeEdge(new_source, new_target);
+  }
+
+  for (const auto& id_edge_pair : other.interlayer_edges()) {
+    const auto& edge = id_edge_pair.second;
+    NodeId new_source = config.getMergedId(edge.source);
+    NodeId new_target = config.getMergedId(edge.target);
+    if (new_source == new_target) {
+      continue;
+    }
+
+    if (config.enforce_parent_constraints) {
+      insertParentEdge(new_source, new_target, edge.info->clone());
+    } else {
+      insertEdge(new_source, new_target, edge.info->clone());
+    }
+  }
+}
+
 Edge* EdgeContainer::find(const EdgeKey& key) const {
   auto iter = edges.find(key);
   if (iter == edges.end()) {
@@ -152,5 +212,4 @@ Edge* EdgeContainer::find(const EdgeKey& key) const {
   return const_cast<SceneGraphEdge*>(&iter->second);
 }
 
-
 }  // namespace spark_dsg
diff --git a/src/scene_graph_layer.cpp b/src/scene_graph_layer.cpp
index 08086bf..52ae87d 100644
--- a/src/scene_graph_layer.cpp
+++ b/src/scene_graph_layer.cpp
@@ -248,23 +248,7 @@ bool SceneGraphLayer::mergeLayer(const SceneGraphLayer& other_layer,
     }
   }
 
-  for (const auto& id_edge_pair : other_layer.edges_.edges) {
-    const auto& edge = id_edge_pair.second;
-    NodeId new_source = config.getMergedId(edge.source);
-    NodeId new_target = config.getMergedId(edge.target);
-    if (new_source == new_target) {
-      continue;
-    }
-
-    // NOTE(nathan) not really necessary but saves the clone call
-    if (hasEdge(new_source, new_target)) {
-      // TODO(nathan) clone attributes
-      continue;
-    }
-
-    insertEdge(new_source, new_target, edge.info->clone());
-  }
-
+  edges_.updateFrom(other_layer.edges_);
   return true;
 }