From 5950ab81b2d1adf36354e3b9e7bc5b16e0a97874 Mon Sep 17 00:00:00 2001
From: jaensen <4954577+jaensen@users.noreply.github.com>
Date: Tue, 30 May 2023 20:07:29 +0200
Subject: [PATCH] add flow-edges from users with trusted tokens to the
 organizations that trust them.

---
 src/safe_db/db.rs | 41 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/safe_db/db.rs b/src/safe_db/db.rs
index e331b96..38a51e2 100644
--- a/src/safe_db/db.rs
+++ b/src/safe_db/db.rs
@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap, HashSet};
 
 use crate::types::{edge::EdgeDB, Address, Edge, Safe, U256};
 
@@ -31,6 +31,45 @@ impl DB {
 
     fn compute_edges(&mut self) {
         let mut edges = vec![];
+
+        // token address -> orga addresses
+        let mut organization_accepted_tokens: HashMap<Address, HashSet<Address>> = HashMap::new();
+
+        // Build a map from token address to orga addresses that accept this token
+        for (_, safe) in &self.safes {
+            for (send_to, percentage) in &safe.limit_percentage {
+                if percentage == &0 {
+                    continue;
+                }
+
+                let receiver_safe = self.safes.get(send_to).unwrap();
+                if receiver_safe.organization {
+                    //println!("user {} can send {} token to orga {}", user, safe.token_address, send_to);
+                    organization_accepted_tokens.entry(safe.token_address).or_default().insert(*send_to);
+                }
+            }
+        }
+
+        // Find all safes that have a non-zero balance of tokens that are accepted by an organization
+        for (user, safe) in &self.safes {
+            for (token, balance) in &safe.balances {
+                if balance == &U256::from(0) {
+                    continue;
+                }
+                organization_accepted_tokens.get(token).map(|organizations| {
+                    for organization in organizations {
+                        // Add the balance as capacity from 'user' to 'organization'
+                        edges.push(Edge {
+                            from: *user,
+                            to: *organization,
+                            token: *token,
+                            capacity: *balance,
+                        });
+                    }
+                });
+            }
+        }
+
         for (user, safe) in &self.safes {
             // trust connections
             for (send_to, percentage) in &safe.limit_percentage {