refactor: AllSenderPeers and ReceiverPeers

This commit is contained in:
Youngjoon Lee 2024-08-25 00:45:26 +09:00
parent 2976a53628
commit 138ae1f5c9
No known key found for this signature in database
GPG Key ID: 167546E2D1712F8C
4 changed files with 88 additions and 41 deletions

View File

@ -15,7 +15,7 @@ use crate::{
};
pub fn run_iteration(paramset: ParamSet, seed: u64, outputs: &mut Outputs) -> f32 {
let (mut mixnodes, sender_peers_list, receiver_peer_conn_idx) = if paramset.random_topology {
let (mut mixnodes, all_sender_peers, receiver_peers) = if paramset.random_topology {
build_random_network(&paramset, seed, outputs)
} else {
build_striped_network(&paramset, seed)
@ -56,9 +56,9 @@ pub fn run_iteration(paramset: ParamSet, seed: u64, outputs: &mut Outputs) -> f3
// All senders emit a message (data or noise) to all of their own adjacent peers.
if all_sent_count < all_sent_count_target {
// For each sender
for (sender_idx, sender_peers) in sender_peers_list.iter().enumerate() {
for (sender_idx, sender_peers) in all_sender_peers.iter() {
if try_probability(&mut data_msg_rng, paramset.sender_data_msg_prob) {
let msg = data_msg_gen.next(sender_idx.try_into().unwrap());
let msg = data_msg_gen.next(sender_idx);
sender_peers.iter().for_each(|peer_id| {
mixnodes
.get_mut(*peer_id as usize)
@ -107,13 +107,13 @@ pub fn run_iteration(paramset: ParamSet, seed: u64, outputs: &mut Outputs) -> f3
}
}
// Record msg to the sequence
let conn_idx = receiver_peer_conn_idx.get(&relayer_id).unwrap();
outputs.add_recv_msg(&msg, *conn_idx as usize);
let conn_idx = receiver_peers.conn_idx(&relayer_id).unwrap();
outputs.add_recv_msg(&msg, conn_idx);
}
Message::Noise => {
// Record noise to the sequence
let conn_idx = receiver_peer_conn_idx.get(&relayer_id).unwrap();
outputs.add_recv_noise(*conn_idx as usize);
let conn_idx = receiver_peers.conn_idx(&relayer_id).unwrap();
outputs.add_recv_noise(conn_idx);
}
}
} else if let Message::Data(msg) = msg {

View File

@ -1,8 +1,10 @@
use std::fmt::Display;
pub type SenderIdx = u8;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct DataMessage {
pub sender: u8,
pub sender: SenderIdx,
pub msg_id: u32,
}
@ -23,7 +25,7 @@ impl DataMessageGenerator {
}
}
pub fn next(&mut self, sender: u8) -> DataMessage {
pub fn next(&mut self, sender: SenderIdx) -> DataMessage {
let msg_id = self.next_msg_ids[sender as usize];
self.next_msg_ids[sender as usize] += 1;
DataMessage { sender, msg_id }

View File

@ -5,7 +5,11 @@ use protocol::{
topology::Topology,
};
use crate::{message::DataMessage, ordercoeff::SequenceWriter};
use crate::{
message::{DataMessage, SenderIdx},
ordercoeff::SequenceWriter,
topology::AllSenderPeers,
};
pub struct Outputs {
closed: bool,
@ -103,18 +107,18 @@ impl Outputs {
writer.add_message(msg);
}
pub fn add_sent_noise(&mut self, sender_idx: usize) {
let writer = &mut self.sent_sequence_writers[sender_idx];
pub fn add_sent_noise(&mut self, sender_idx: SenderIdx) {
let writer = &mut self.sent_sequence_writers[sender_idx as usize];
writer.add_noise();
}
pub fn add_recv_msg(&mut self, msg: &DataMessage, conn_idx: usize) {
let writer = &mut self.recv_sequence_writers[conn_idx];
pub fn add_recv_msg(&mut self, msg: &DataMessage, conn_idx: u16) {
let writer = &mut self.recv_sequence_writers[conn_idx as usize];
writer.add_message(msg);
}
pub fn add_recv_noise(&mut self, conn_idx: usize) {
let writer = &mut self.recv_sequence_writers[conn_idx];
pub fn add_recv_noise(&mut self, conn_idx: u16) {
let writer = &mut self.recv_sequence_writers[conn_idx as usize];
writer.add_noise();
}
@ -151,7 +155,7 @@ impl Outputs {
pub fn write_topology(
&self,
topology: &Topology,
sender_peers_list: &[Vec<NodeId>],
all_sender_peers: &AllSenderPeers,
receiver_peers: &[NodeId],
) {
let mut writer = csv::Writer::from_path(&self.topology_path).unwrap();
@ -176,7 +180,7 @@ impl Outputs {
}
// Write peers of senders
for (sender_idx, peers) in sender_peers_list.iter().enumerate() {
for (sender_idx, peers) in all_sender_peers.iter() {
writer
.write_record(&[
format!("sender-{}", sender_idx),

View File

@ -8,14 +8,14 @@ use protocol::{
use rand::{rngs::StdRng, seq::SliceRandom, RngCore, SeedableRng};
use rustc_hash::FxHashMap;
use crate::{outputs::Outputs, paramset::ParamSet};
use crate::{message::SenderIdx, outputs::Outputs, paramset::ParamSet};
pub const RECEIVER_NODE_ID: NodeId = NodeId::MAX;
pub fn build_striped_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq + Hash>(
paramset: &ParamSet,
seed: u64,
) -> (Vec<Node<M>>, Vec<Vec<NodeId>>, FxHashMap<NodeId, u16>) {
) -> (Vec<Node<M>>, AllSenderPeers, ReceiverPeers) {
assert!(!paramset.random_topology);
let mut next_node_id: NodeId = 0;
let mut queue_seed_rng = StdRng::seed_from_u64(seed);
@ -43,7 +43,7 @@ pub fn build_striped_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq
}
// Connect mix nodes
let mut receiver_peer_conn_idx: FxHashMap<NodeId, u16> = FxHashMap::default();
let mut receiver_peers = ReceiverPeers::new();
for path in paths.iter() {
for (i, id) in path.iter().enumerate() {
if i != path.len() - 1 {
@ -56,24 +56,28 @@ pub fn build_striped_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq
assert_eq!(mixnode.id, *id);
mixnode.connect(RECEIVER_NODE_ID);
receiver_peer_conn_idx
.insert(*id, receiver_peer_conn_idx.len().try_into().unwrap());
receiver_peers.add(*id, receiver_peers.len().try_into().unwrap());
}
}
}
let sender_peers_list: Vec<Vec<NodeId>> =
vec![
paths.iter().map(|path| *path.first().unwrap()).collect();
paramset.num_senders as usize
];
(mixnodes, sender_peers_list, receiver_peer_conn_idx)
let mut all_sender_peers = AllSenderPeers::new(paramset.num_senders);
let sender_peers = paths
.iter()
.map(|path| *path.first().unwrap())
.collect::<Vec<_>>();
(0..paramset.num_senders).for_each(|_| {
all_sender_peers.add(sender_peers.clone());
});
(mixnodes, all_sender_peers, receiver_peers)
}
pub fn build_random_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq + Hash>(
paramset: &ParamSet,
seed: u64,
outputs: &mut Outputs,
) -> (Vec<Node<M>>, Vec<Vec<NodeId>>, FxHashMap<NodeId, u16>) {
) -> (Vec<Node<M>>, AllSenderPeers, ReceiverPeers) {
assert!(paramset.random_topology);
// Init mix nodes
let mut queue_seed_rng = StdRng::seed_from_u64(seed);
@ -95,7 +99,7 @@ pub fn build_random_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq +
let mut peers_rng = StdRng::seed_from_u64(seed);
let mut candidates: Vec<NodeId> = mixnodes.iter().map(|mixnode| mixnode.id).collect();
assert!(candidates.len() >= paramset.peering_degree as usize);
let mut sender_peers_list: Vec<Vec<NodeId>> = Vec::with_capacity(paramset.num_senders as usize);
let mut all_sender_peers = AllSenderPeers::new(paramset.num_senders);
for _ in 0..paramset.num_senders {
candidates.as_mut_slice().shuffle(&mut peers_rng);
let mut peers: Vec<NodeId> = candidates
@ -104,15 +108,15 @@ pub fn build_random_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq +
.take(paramset.peering_degree as usize)
.collect();
peers.sort();
sender_peers_list.push(peers);
all_sender_peers.add(peers);
}
candidates.as_mut_slice().shuffle(&mut peers_rng);
let mut receiver_peers: Vec<NodeId> = candidates
let mut receiver_peer_ids: Vec<NodeId> = candidates
.iter()
.cloned()
.take(paramset.peering_degree as usize)
.collect();
receiver_peers.sort();
receiver_peer_ids.sort();
// Connect mix nodes
let topology = build_topology(
@ -129,18 +133,55 @@ pub fn build_random_network<M: 'static + Debug + Copy + Clone + PartialEq + Eq +
}
// Connect the selected mix nodes with the receiver
//
// peer_id -> conn_idx
let mut receiver_peer_conn_idx: FxHashMap<NodeId, u16> = FxHashMap::default();
for (conn_idx, mixnode_id) in receiver_peers.iter().enumerate() {
let mut receiver_peers = ReceiverPeers::new();
for (conn_idx, mixnode_id) in receiver_peer_ids.iter().enumerate() {
let mixnode = mixnodes.get_mut(*mixnode_id as usize).unwrap();
assert_eq!(mixnode.id, *mixnode_id);
mixnode.connect(RECEIVER_NODE_ID);
receiver_peer_conn_idx.insert(*mixnode_id, conn_idx.try_into().unwrap());
receiver_peers.add(*mixnode_id, conn_idx.try_into().unwrap());
}
outputs.write_topology(&topology, &sender_peers_list, &receiver_peers);
outputs.write_topology(&topology, &all_sender_peers, &receiver_peer_ids);
(mixnodes, sender_peers_list, receiver_peer_conn_idx)
(mixnodes, all_sender_peers, receiver_peers)
}
pub struct AllSenderPeers(Vec<Vec<NodeId>>);
impl AllSenderPeers {
fn new(num_senders: u8) -> Self {
Self(Vec::with_capacity(num_senders as usize))
}
fn add(&mut self, peers: Vec<NodeId>) {
self.0.push(peers)
}
pub fn iter(&self) -> impl Iterator<Item = (SenderIdx, &Vec<NodeId>)> {
self.0
.iter()
.enumerate()
.map(|(idx, v)| (idx.try_into().unwrap(), v))
}
}
pub struct ReceiverPeers(FxHashMap<NodeId, u16>);
impl ReceiverPeers {
fn new() -> Self {
ReceiverPeers(FxHashMap::default())
}
fn add(&mut self, peer_id: NodeId, conn_idx: u16) {
self.0.insert(peer_id, conn_idx);
}
pub fn conn_idx(&self, node_id: &NodeId) -> Option<u16> {
self.0.get(node_id).cloned()
}
fn len(&self) -> usize {
self.0.len()
}
}