1use crate::time::ReplicaId;
  2use ctor::ctor;
  3use rand::Rng;
  4use std::{
  5    collections::BTreeMap,
  6    path::{Path, PathBuf},
  7};
  8use tempdir::TempDir;
  9
 10#[ctor]
 11fn init_logger() {
 12    env_logger::init();
 13}
 14
 15#[derive(Clone)]
 16struct Envelope<T: Clone> {
 17    message: T,
 18    sender: ReplicaId,
 19}
 20
 21pub(crate) struct Network<T: Clone> {
 22    inboxes: BTreeMap<ReplicaId, Vec<Envelope<T>>>,
 23    all_messages: Vec<T>,
 24}
 25
 26impl<T: Clone> Network<T> {
 27    pub fn new() -> Self {
 28        Network {
 29            inboxes: BTreeMap::new(),
 30            all_messages: Vec::new(),
 31        }
 32    }
 33
 34    pub fn add_peer(&mut self, id: ReplicaId) {
 35        self.inboxes.insert(id, Vec::new());
 36    }
 37
 38    pub fn is_idle(&self) -> bool {
 39        self.inboxes.values().all(|i| i.is_empty())
 40    }
 41
 42    pub fn broadcast<R>(&mut self, sender: ReplicaId, messages: Vec<T>, rng: &mut R)
 43    where
 44        R: Rng,
 45    {
 46        for (replica, inbox) in self.inboxes.iter_mut() {
 47            if *replica != sender {
 48                for message in &messages {
 49                    let min_index = inbox
 50                        .iter()
 51                        .enumerate()
 52                        .rev()
 53                        .find_map(|(index, envelope)| {
 54                            if sender == envelope.sender {
 55                                Some(index + 1)
 56                            } else {
 57                                None
 58                            }
 59                        })
 60                        .unwrap_or(0);
 61
 62                    // Insert one or more duplicates of this message *after* the previous
 63                    // message delivered by this replica.
 64                    for _ in 0..rng.gen_range(1..4) {
 65                        let insertion_index = rng.gen_range(min_index..inbox.len() + 1);
 66                        inbox.insert(
 67                            insertion_index,
 68                            Envelope {
 69                                message: message.clone(),
 70                                sender,
 71                            },
 72                        );
 73                    }
 74                }
 75            }
 76        }
 77        self.all_messages.extend(messages);
 78    }
 79
 80    pub fn has_unreceived(&self, receiver: ReplicaId) -> bool {
 81        !self.inboxes[&receiver].is_empty()
 82    }
 83
 84    pub fn receive<R>(&mut self, receiver: ReplicaId, rng: &mut R) -> Vec<T>
 85    where
 86        R: Rng,
 87    {
 88        let inbox = self.inboxes.get_mut(&receiver).unwrap();
 89        let count = rng.gen_range(0..inbox.len() + 1);
 90        inbox
 91            .drain(0..count)
 92            .map(|envelope| envelope.message)
 93            .collect()
 94    }
 95}
 96
 97pub fn sample_text(rows: usize, cols: usize) -> String {
 98    let mut text = String::new();
 99    for row in 0..rows {
100        let c: char = ('a' as u32 + row as u32) as u8 as char;
101        let mut line = c.to_string().repeat(cols);
102        if row < rows - 1 {
103            line.push('\n');
104        }
105        text += &line;
106    }
107    text
108}
109
110pub fn temp_tree(tree: serde_json::Value) -> TempDir {
111    let dir = TempDir::new("").unwrap();
112    write_tree(dir.path(), tree);
113    dir
114}
115
116fn write_tree(path: &Path, tree: serde_json::Value) {
117    use serde_json::Value;
118    use std::fs;
119
120    if let Value::Object(map) = tree {
121        for (name, contents) in map {
122            let mut path = PathBuf::from(path);
123            path.push(name);
124            match contents {
125                Value::Object(_) => {
126                    fs::create_dir(&path).unwrap();
127                    write_tree(&path, contents);
128                }
129                Value::Null => {
130                    fs::create_dir(&path).unwrap();
131                }
132                Value::String(contents) => {
133                    fs::write(&path, contents).unwrap();
134                }
135                _ => {
136                    panic!("JSON object must contain only objects, strings, or null");
137                }
138            }
139        }
140    } else {
141        panic!("You must pass a JSON object to this helper")
142    }
143}