From 18424a2fc0ca6c5c054ab7a769f574ab2e062222 Mon Sep 17 00:00:00 2001
From: Tarek <tareknaser360@gmail.com>
Date: Fri, 31 May 2024 16:35:04 +0300
Subject: [PATCH] feat: add a method to visualize the network graph

Signed-off-by: Tarek <tareknaser360@gmail.com>
---
 Cargo.lock               | 39 +++++++++++++++++++++++++++++++++++++++
 barq-common/Cargo.toml   |  3 +++
 barq-common/src/graph.rs | 24 ++++++++++++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index fa8d49e..7fd1746 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -13,6 +13,7 @@ name = "barq-common"
 version = "0.0.1"
 dependencies = [
  "log",
+ "petgraph",
  "serde",
  "serde_json",
 ]
@@ -100,6 +101,34 @@ version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fb4a24b1aaf0fd0ce8b45161144d6f42cd91677fd5940fd431183eb023b3a2b8"
 
+[[package]]
+name = "equivalent"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
+
+[[package]]
+name = "fixedbitset"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
+
+[[package]]
+name = "hashbrown"
+version = "0.14.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
+
+[[package]]
+name = "indexmap"
+version = "2.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26"
+dependencies = [
+ "equivalent",
+ "hashbrown",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.11"
@@ -118,6 +147,16 @@ version = "0.4.21"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c"
 
+[[package]]
+name = "petgraph"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db"
+dependencies = [
+ "fixedbitset",
+ "indexmap",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.83"
diff --git a/barq-common/Cargo.toml b/barq-common/Cargo.toml
index 542732a..a7463cd 100644
--- a/barq-common/Cargo.toml
+++ b/barq-common/Cargo.toml
@@ -9,3 +9,6 @@ edition = "2021"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 log = "0.4"
+
+# Dependencies to visualize the graph
+petgraph = "0.6"
diff --git a/barq-common/src/graph.rs b/barq-common/src/graph.rs
index ac3e18c..6d73332 100644
--- a/barq-common/src/graph.rs
+++ b/barq-common/src/graph.rs
@@ -1,6 +1,11 @@
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
+use petgraph::{
+    dot::{Config, Dot},
+    graph::{DiGraph, NodeIndex},
+};
+
 /// Represents a node in the network graph.
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct Node {
@@ -113,6 +118,25 @@ impl NetworkGraph {
 
     // TODO: Add methods for updating nodes and edges.
     // TODO: Add methods for removing nodes and edges.
+
+    /// Returns a DOT representation of the network graph.
+    pub fn to_dot(&self) -> String {
+        let mut graph = DiGraph::new();
+        let mut node_indices: HashMap<String, NodeIndex> = HashMap::new();
+
+        for node in self.nodes.values() {
+            let node_index = graph.add_node(node.id.clone());
+            node_indices.insert(node.id.clone(), node_index);
+        }
+
+        for edge in self.edges.values() {
+            let source_index = node_indices.get(&edge.source).unwrap();
+            let destination_index = node_indices.get(&edge.destination).unwrap();
+            graph.add_edge(*source_index, *destination_index, edge.capacity);
+        }
+
+        format!("{:?}", Dot::with_config(&graph, &[Config::EdgeNoLabel]))
+    }
 }
 
 #[cfg(test)]