rust.clogsim/src/main.rs

extern crate sha2;

use sha2::{Sha256, Digest};

use std::io;
use std::io::prelude::*;

use std::collections::VecDeque;
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;

use std::time::{SystemTime, UNIX_EPOCH};

use std::fmt;

const INTRO: &'static str = "This is a small program that demonstrates some concepts of clog.\n\
                             To begin, type \"help\".";
const HELP: &'static str = "Available commands:\n\
                         \x20   help - Prints this help text\n\
                         \x20   status [server] - Prints queues, buffers, details, etc\n\
                         \x20   new <server> - Creates a new server\n\
                         \x20   kill <server> - Removes a server\n\
                         \x20   send <server> <message> - Sends a message to a server\n\
                         \x20   recv <from-server> <to-server> - Sends a message between servers\n\
                         \x20   quit - Exits this program\n\
                         ";

#[derive(PartialEq, Eq, Hash, Ord, PartialOrd, Clone, Default, Debug)]
struct Hash {
    ty: String,
    depth: u64,
    hash: String,
}

impl Hash {
    fn of(s: &str, d: u64) -> Hash {
        Hash {
            ty: "sha256".into(),
            depth: d + 1,
            hash: format!("{:x}", Sha256::digest(s.as_bytes())),
        }
    }
}

impl fmt::Display for Hash {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}={}/{}", self.ty, self.depth, self.hash)
    }
}

#[derive(Default)]
struct ServerData {
    hashes: BTreeSet<Hash>,
    message_queue: HashMap<String, VecDeque<(BTreeSet<Hash>, String)>>, // (source) server name -> messages
    users: BTreeSet<String>,
}

impl ServerData {
    fn on_msg(&mut self, (hashes, msg): (BTreeSet<Hash>, String)) {
        let new_hashes = self.hashes.difference(&hashes).cloned().collect();
        self.hashes = new_hashes;
        let mut h_msg = format!("hash=");
        let mut depth: u64 = 0;
        for hash in hashes {
            h_msg.push_str(&hash.to_string());
            depth = depth.max(hash.depth);
        }
        h_msg.push_str(&format!(" {}", msg));
        let next_hash: Hash = Hash::of(&h_msg, depth);
        self.hashes.insert(next_hash);
    }
}

#[derive(Default)]
struct Game {
    servers: BTreeMap<String, ServerData>,
    users: BTreeSet<String>,
}

impl Game {
    fn new(&mut self, s: &str) {
        if let Some(server_name) = singular(s.trim().split_whitespace()) {
            let mut server_data = ServerData::default();
            // add all known hashes, skip duplicates
            server_data.hashes.extend(self.servers.values().flat_map(|serv| &serv.hashes).cloned());
            self.servers.insert(server_name.to_owned(), server_data);
            println!("Added server: {}", server_name);
        } else {
            println!("Syntax: new <server>\nSynopsis: Creates a new server");
        }
    }
    fn print_status(&self, s: &str) {
        if let Some(server_name) = singular(s.trim().split_whitespace()) {
            if let Some(server) = self.servers.get(server_name) {
                println!("{}:", server_name);
                println!("    Hashes:");
                for hash in &server.hashes {
                    println!("        {}", hash);
                }
                println!("    Messages:");
                for msg in server.message_queue.iter().flat_map(|(k, v)| v.iter().map(move |v| format_msg(k, v))) {
                    println!("        {}", msg);
                }
            } else {
                println!("Unknown server: {}", server_name);
            }
        } else {
            for server in self.servers.keys() {
                self.print_status(server);
            }
            if self.servers.is_empty() {
                println!("No servers");
            }
        }
    }
    fn kill(&mut self, s: &str) {
        if let Some(server_name) = singular(s.trim().split_whitespace()) {
            if let Some(_server) = self.servers.remove(server_name) {
                println!("Removed server: {}", server_name);
            } else {
                println!("Unknown server: {}", server_name);
            }
        } else {
            println!("Syntax: kill <server>\nSynopsis: Removes a server");
        }
    }
    fn send(&mut self, s: &str) {
        let mut it = s.splitn(2, " ");
        if let (Some(server_name), Some(msg)) = (it.next().and_then(|sn| singular(sn.trim().split_whitespace())), it.next().map(|m| m.trim())) {
            let msg = format!("[{}] {}", SystemTime::now().duration_since(UNIX_EPOCH).expect("time before unix epoch").as_secs(), msg);
            if let Some(hashes) = self.servers.get(server_name).map(|serv| serv.hashes.clone()) {
                for server in self.servers.values_mut() {
                    server.message_queue.entry(server_name.to_owned()).or_default().push_back((hashes.clone(), msg.to_owned()));
                }
            }
            if let Some(server) = self.servers.get_mut(server_name) {
                if let Some(msg) = server.message_queue.get_mut(server_name).and_then(|msgs| msgs.pop_front()) {
                    server.on_msg(msg);
                } else {
                    panic!("Somehow got a None in a place you're not supposed to get a None");
                }
            } else {
                println!("Unknown server: {}", server_name);
            }
        } else {
            println!("Syntax: send <server> <message>\nSynopsis: Sends a message to a server");
        }
    }
    fn recv(&mut self, s: &str) {
        let mut it = s.splitn(2, " ");
        if let (Some(from), Some(to)) = (it.next().and_then(|sn| singular(sn.trim().split_whitespace())), it.next().and_then(|sn| singular(sn.trim().split_whitespace()))) {
            if let Some(server) = self.servers.get_mut(to) {
                if let Some(msg) = server.message_queue.get_mut(from).and_then(|msgs| msgs.pop_front()) {
                    server.on_msg(msg);
                } else {
                    println!("No message from server: {}", from);
                }
            } else {
                println!("Unknown server: {}", to);
            }
        } else {
            println!("Syntax: recv <from-server> <to-server>\nSynopsis: Sends a message between servers");
        }
    }
}

fn format_msg(serv_name: &str, msg: &(BTreeSet<Hash>, String)) -> String {
    let mut s = format!("[{}] hash=", serv_name);
    for hash in msg.0.iter() {
        s.push_str(&format!("{},", hash));
    }
    s.push_str(&format!(" {}", msg.1));
    s
}

fn singular<T, I: Iterator<Item=T>>(mut it: I) -> Option<T> {
    let res = it.next();
    res.filter(|_| it.next().is_none())
}

fn main() {
    println!("{}", INTRO);
    let stdin = io::stdin();
    let mut game = Game::default();
    loop {
        print!("> ");
        let _ = io::stdout().flush();
        let mut line = String::new();
        if stdin.read_line(&mut line).is_ok() {
            let mut it = line.trim().splitn(2, " ");
            if let Some(s) = it.next() { match s {
                "help" => println!("{}", HELP),
                "quit" => return,
                "status" => game.print_status(it.next().unwrap_or("")),
                "new" => game.new(it.next().unwrap_or("")),
                "kill" => game.kill(it.next().unwrap_or("")),
                "send" => game.send(it.next().unwrap_or("")),
                "recv" => game.recv(it.next().unwrap_or("")),
                s => println!("Unknown command: {}", s),
            } } else { println!("No input. Try \"help\"."); }
        }
    }
}