From d15c421a3aeaa5727df98d58e940f97c8063095a Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Sun, 25 Aug 2024 04:30:50 +0900
Subject: [PATCH] fix: receiver seq counts

---
 mixnet-rs/ordering/src/iteration.rs |  2 +-
 mixnet-rs/ordering/src/outputs.rs   |  8 ++++----
 mixnet-rs/ordering/src/paramset.rs  |  8 ++++++++
 mixnet-rs/ordering/src/topology.rs  | 10 +++++-----
 4 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/mixnet-rs/ordering/src/iteration.rs b/mixnet-rs/ordering/src/iteration.rs
index 9056732..31e5d57 100644
--- a/mixnet-rs/ordering/src/iteration.rs
+++ b/mixnet-rs/ordering/src/iteration.rs
@@ -34,7 +34,7 @@ impl Iteration {
             (0..self.paramset.num_senders)
                 .map(|sender_idx| format!("{dir}/sent_seq_{sender_idx}__WIP__.csv"))
                 .collect(),
-            (0..self.paramset.peering_degree)
+            (0..self.paramset.num_receiver_connections())
                 .map(|conn_idx| format!("{dir}/recv_seq_{conn_idx}__WIP__.csv"))
                 .collect(),
             format!("{dir}/data_msg_counts__WIP__.csv"),
diff --git a/mixnet-rs/ordering/src/outputs.rs b/mixnet-rs/ordering/src/outputs.rs
index 9f1dccf..06c34da 100644
--- a/mixnet-rs/ordering/src/outputs.rs
+++ b/mixnet-rs/ordering/src/outputs.rs
@@ -112,13 +112,13 @@ impl Outputs {
         writer.add_noise();
     }
 
-    pub fn add_recv_msg(&mut self, msg: &DataMessage, conn_idx: u16) {
-        let writer = &mut self.recv_sequence_writers[conn_idx as usize];
+    pub fn add_recv_msg(&mut self, msg: &DataMessage, conn_idx: usize) {
+        let writer = &mut self.recv_sequence_writers[conn_idx];
         writer.add_message(msg);
     }
 
-    pub fn add_recv_noise(&mut self, conn_idx: u16) {
-        let writer = &mut self.recv_sequence_writers[conn_idx as usize];
+    pub fn add_recv_noise(&mut self, conn_idx: usize) {
+        let writer = &mut self.recv_sequence_writers[conn_idx];
         writer.add_noise();
     }
 
diff --git a/mixnet-rs/ordering/src/paramset.rs b/mixnet-rs/ordering/src/paramset.rs
index e91e35c..8a08c2e 100644
--- a/mixnet-rs/ordering/src/paramset.rs
+++ b/mixnet-rs/ordering/src/paramset.rs
@@ -179,6 +179,14 @@ impl ParamSet {
         paramsets
     }
 
+    pub fn num_receiver_connections(&self) -> usize {
+        if self.random_topology {
+            self.peering_degree as usize
+        } else {
+            self.num_paths as usize
+        }
+    }
+
     pub fn as_csv_record(&self) -> Vec<String> {
         vec![
             self.id.to_string(),
diff --git a/mixnet-rs/ordering/src/topology.rs b/mixnet-rs/ordering/src/topology.rs
index 7d81e96..8a78f39 100644
--- a/mixnet-rs/ordering/src/topology.rs
+++ b/mixnet-rs/ordering/src/topology.rs
@@ -56,7 +56,7 @@ pub fn build_striped_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq
                 assert_eq!(mixnode.id, *id);
                 mixnode.connect(RECEIVER_NODE_ID);
 
-                receiver_peers.add(*id, receiver_peers.len().try_into().unwrap());
+                receiver_peers.add(*id, receiver_peers.len());
             }
         }
     }
@@ -139,7 +139,7 @@ pub fn build_random_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq +
         assert_eq!(mixnode.id, *mixnode_id);
         mixnode.connect(RECEIVER_NODE_ID);
 
-        receiver_peers.add(*mixnode_id, conn_idx.try_into().unwrap());
+        receiver_peers.add(*mixnode_id, conn_idx);
     }
 
     outputs.write_topology(&topology, &all_sender_peers, &receiver_peer_ids);
@@ -166,18 +166,18 @@ impl AllSenderPeers {
     }
 }
 
-pub struct ReceiverPeers(FxHashMap<NodeId, u16>);
+pub struct ReceiverPeers(FxHashMap<NodeId, usize>);
 
 impl ReceiverPeers {
     fn new() -> Self {
         ReceiverPeers(FxHashMap::default())
     }
 
-    fn add(&mut self, peer_id: NodeId, conn_idx: u16) {
+    fn add(&mut self, peer_id: NodeId, conn_idx: usize) {
         self.0.insert(peer_id, conn_idx);
     }
 
-    pub fn conn_idx(&self, node_id: &NodeId) -> Option<u16> {
+    pub fn conn_idx(&self, node_id: &NodeId) -> Option<usize> {
         self.0.get(node_id).cloned()
     }