From 794266980cd1160a359a084626673fc790da0304 Mon Sep 17 00:00:00 2001
From: Sarah Jamie Lewis <sarah@openprivacy.ca>
Date: Wed, 3 Feb 2021 14:46:52 -0800
Subject: [PATCH] Nicer visualizations

---
 src/oracle.rs | 49 +++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 45 insertions(+), 4 deletions(-)

diff --git a/src/oracle.rs b/src/oracle.rs
index cebe4e3..44fb658 100644
--- a/src/oracle.rs
+++ b/src/oracle.rs
@@ -1,6 +1,7 @@
 // todo would be neat to just make this a subscriber of the tracing log, but there doesn't
 // seem to be a nice way for that API to do what we want.. so until then...
 
+use hashbrown::HashMap;
 use rand::{thread_rng, Rng};
 use std::fs::File;
 use std::io::Write;
@@ -46,12 +47,12 @@ impl Oracle {
         write!(output, "digraph {{\n");
         write!(
             output,
-            r#"K=2.5;
+            r#"K=2;
                         repulsiveforce=0.1;
                         overlap=true;
                         splines = true;
                         dpi=400;
-                        penwidth = 0.8;
+                        penwidth = 1;
                         "#
         );
 
@@ -62,15 +63,55 @@ impl Oracle {
             writeln!(output, "\"{}\" [shape=point, color=\"#{:x}{:x}{:x}\"]", party, r, g, b);
         }
 
+        let mut real_connection_map: HashMap<(String, String), u64> = HashMap::new();
+        let mut entangled_connection_map: HashMap<(String, String), u64> = HashMap::new();
+        let mut max_conn = 1;
         for event in self.actual_events.iter() {
-            writeln!(output, "\"{}\" -> \"{}\" [arrowhead=none]", event.sender, event.intended_receiver);
+            let key = (event.sender.clone(), event.intended_receiver.clone());
+            if real_connection_map.contains_key(&key) {
+                *real_connection_map.get_mut(&key).unwrap() += 1;
+                if real_connection_map[&key] > max_conn {
+                    max_conn = real_connection_map[&key];
+                }
+            } else {
+                real_connection_map.insert(key, 1);
+            }
             match &event.entangled_receiver {
                 Some(entangled_receiver) => {
-                    writeln!(output, "\"{}\" -> \"{}\" [arrowhead=none,style=dashed]", event.sender, entangled_receiver);
+                    let key = (event.sender.clone(), entangled_receiver.clone());
+                    if entangled_connection_map.contains_key(&key) {
+                        *entangled_connection_map.get_mut(&key).unwrap() += 1;
+                        if entangled_connection_map[&key] > max_conn {
+                            max_conn = entangled_connection_map[&key];
+                        }
+                    } else {
+                        entangled_connection_map.insert(key, 1);
+                    }
                 }
                 _ => {}
             };
         }
+
+        for ((sender, receiver), size) in real_connection_map.iter() {
+            let normalized = (*size as f64 / max_conn as f64);
+            let transparency = (normalized * 255.0) as u8;
+            let penwidth = (normalized * 2.0) as f64;
+            writeln!(
+                output,
+                "\"{}\" -> \"{}\" [arrowhead=none, penwidth={}, color=\"#000000{:x}\"]",
+                sender, receiver, penwidth, transparency
+            );
+        }
+        for ((sender, receiver), size) in entangled_connection_map.iter() {
+            let normalized = (*size as f64 / max_conn as f64);
+            let transparency = (normalized * 255.0) as u8;
+            let penwidth = (normalized * 2.0) as f64;
+            writeln!(
+                output,
+                "\"{}\" -> \"{}\" [arrowhead=none, style=dashed, penwidth={}, color=\"#000000{:x}\"]",
+                sender, receiver, penwidth, transparency
+            );
+        }
         write!(output, "}}");
     }
 }