<commit_before>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::{ChatMap, ClientInfo};
use message::Message;

extern crate time;

pub fn create_client(stream: TcpStream,
                     sender_to_router: Sender<Message>,
                     chat_map: &Arc<Mutex<ChatMap>>) {

    println!("New client");
    let mut stream = stream;

    stream.write(&"Welcome to Smazy\n".to_string().into_bytes()).unwrap();
    stream.write(&"Please enter a username:\n".to_string().into_bytes()).unwrap();
    stream.flush().unwrap();

    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut username = "".to_string();

    match reader.read_line(&mut username) {
        Ok(n) => {
            let username = username.trim().to_string();
            let (sender_to_client, receiver_from_router) = channel();

            {
                let mut guard = chat_map.lock().unwrap();
                if n > 0 && !guard.contains_key(&username) {
                    let client_info = ClientInfo{
                        partner: None,
                        sender_to_client: sender_to_client
                    };
                    guard.insert(username.to_string(), client_info);
                    println!("New user: {}", username.to_string());
                }
            }

            let partner = choose_chat_partner(stream.try_clone().unwrap(),
                username.to_string(), chat_map);

            chat(stream, username, partner, sender_to_router, receiver_from_router);
        },

        Err(e) => println!("{}", e)
    }
}

fn choose_chat_partner(stream: TcpStream, username: String,
    chat_map: &Arc<Mutex<ChatMap>>) -> String {

    let mut stream = stream;

    let show_available_msg = "Here are the users available to chat:\n".to_string();
    stream.write(&show_available_msg.into_bytes()).unwrap();

    {
        let guard = chat_map.lock().unwrap();
        for (name, client_info) in guard.iter() {
            if &name[..] != &username[..] && client_info.partner == None {
                stream.write(&name.to_string().into_bytes()).unwrap();
                stream.write(&"\n".to_string().into_bytes()).unwrap();
            }
        }
    }

    let select_msg = "Select who you want to chat with:\n".to_string();
    stream.write(&select_msg.into_bytes()).unwrap();

    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut partner = "".to_string();

    match reader.read_line(&mut partner) {
        Ok(_) => {
            partner = partner.trim().to_string();

            if !try_select_partner(chat_map, username.to_string(), partner.to_string()) {
                return choose_chat_partner(stream, username, chat_map);
            }
        }
        Err(e) => println!("{}", e)
    }

    partner
}

fn try_select_partner(chat_map: &Arc<Mutex<ChatMap>>,
                      username: String,
                      partner: String) -> bool {

    let success: bool;
    {
        let mut guard = chat_map.lock().unwrap();

        match guard.get_mut(&partner) {
            Some(clientinfo) => {
                if &partner[..] != &username[..]
                    && (clientinfo.partner == None
                        || clientinfo.partner == Some(username.to_string())) {
                    clientinfo.partner = Some(username.to_string());
                    success = true;
                } else {
                    success = false;
                }
            },
            None => success = false
        }
    }

    if success {
        let mut guard = chat_map.lock().unwrap();
        println!("Assigning partner {} to {}", partner, username);
        guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
    }
    success
}

// TODO: Receiving a message while in the middle of typing a message
//       inserts received message into middle of your message
fn chat(stream: TcpStream,
        username: String,
        partner: String,
        sender_to_router: Sender<Message>,
        receiver_from_router: Receiver<Message>) {

    // Send messages
    {
        let username = username.clone();
        let reader = BufReader::new(stream.try_clone().unwrap());
        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);

                let msg = Message::new(time::now().asctime().to_string(),
                                       username.to_string(),
                                       partner.to_string(),
                                       line.to_string());

                sender_to_router.send(msg).unwrap();
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        thread::spawn(move|| {
            loop {
                let stream = stream.try_clone().unwrap();
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}",
                            &username, msg.content());

                        receive_message(stream, msg);
                    },

                    Err(TryRecvError::Empty) => continue,

                    Err(TryRecvError::Disconnected) =>
                        panic!("User {} disconnected from router", &username)
                }
            }
        });
    }
}

fn receive_message(stream: TcpStream, msg: Message) {
    let mut stream = stream;
    let output = msg.from().to_string() + ": " + &msg.content()[..] + "\n";
    stream.write(&output.into_bytes()).unwrap();
}

#[cfg(test)]
mod client_tests {

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::channel;
    use chatmap::{ChatMap, ClientInfo};
    use super::try_select_partner;

    #[test]
    fn try_select_partner_test_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: b already has a partner
        assert!(!try_select_partner(&cm, "c".to_string(), "b".to_string()));

        // Fail: Can't choose self as partner
        assert!(!try_select_partner(&cm, "c".to_string(), "c".to_string()));

        // Success: b can choose a when they are already partners
        //   (This allows someone to confirm a chat request)
        assert!(try_select_partner(&cm, "b".to_string(), "a".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));

        // Fail: c already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "c".to_string()));

        // Fail: d already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "d".to_string()));
    }

    fn fixture() -> ChatMap {
        let mut cm = ChatMap::new();
        let (sender_to_a, _) = channel();
        cm.insert("a".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_a
        });
        let (sender_to_b, _) = channel();
        cm.insert("b".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_b
        });
        let (sender_to_c, _) = channel();
        cm.insert("c".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_c
        });
        let (sender_to_d, _) = channel();
        cm.insert("d".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_d
        });
        let (sender_to_e, _) = channel();
        cm.insert("e".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_e
        });
        cm
    }
}<commit_msg>Separated tests for clarity<commit_after>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::{ChatMap, ClientInfo};
use message::Message;

extern crate time;

pub fn create_client(stream: TcpStream,
                     sender_to_router: Sender<Message>,
                     chat_map: &Arc<Mutex<ChatMap>>) {

    println!("New client");
    let mut stream = stream;

    stream.write(&"Welcome to Smazy\n".to_string().into_bytes()).unwrap();
    stream.write(&"Please enter a username:\n".to_string().into_bytes()).unwrap();
    stream.flush().unwrap();

    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut username = "".to_string();

    match reader.read_line(&mut username) {
        Ok(n) => {
            let username = username.trim().to_string();
            let (sender_to_client, receiver_from_router) = channel();

            {
                let mut guard = chat_map.lock().unwrap();
                if n > 0 && !guard.contains_key(&username) {
                    let client_info = ClientInfo{
                        partner: None,
                        sender_to_client: sender_to_client
                    };
                    guard.insert(username.to_string(), client_info);
                    println!("New user: {}", username.to_string());
                }
            }

            let partner = choose_chat_partner(stream.try_clone().unwrap(),
                username.to_string(), chat_map);

            chat(stream, username, partner, sender_to_router, receiver_from_router);
        },

        Err(e) => println!("{}", e)
    }
}

fn choose_chat_partner(stream: TcpStream, username: String,
    chat_map: &Arc<Mutex<ChatMap>>) -> String {

    let mut stream = stream;

    let show_available_msg = "Here are the users available to chat:\n".to_string();
    stream.write(&show_available_msg.into_bytes()).unwrap();

    {
        let guard = chat_map.lock().unwrap();
        for (name, client_info) in guard.iter() {
            if &name[..] != &username[..] && client_info.partner == None {
                stream.write(&name.to_string().into_bytes()).unwrap();
                stream.write(&"\n".to_string().into_bytes()).unwrap();
            }
        }
    }

    let select_msg = "Select who you want to chat with:\n".to_string();
    stream.write(&select_msg.into_bytes()).unwrap();

    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut partner = "".to_string();

    match reader.read_line(&mut partner) {
        Ok(_) => {
            partner = partner.trim().to_string();

            if !try_select_partner(chat_map, username.to_string(), partner.to_string()) {
                return choose_chat_partner(stream, username, chat_map);
            }
        }
        Err(e) => println!("{}", e)
    }

    partner
}

fn try_select_partner(chat_map: &Arc<Mutex<ChatMap>>,
                      username: String,
                      partner: String) -> bool {

    let success: bool;
    {
        let mut guard = chat_map.lock().unwrap();

        match guard.get_mut(&partner) {
            Some(clientinfo) => {
                if &partner[..] != &username[..]
                    && (clientinfo.partner == None
                        || clientinfo.partner == Some(username.to_string())) {
                    clientinfo.partner = Some(username.to_string());
                    success = true;
                } else {
                    success = false;
                }
            },
            None => success = false
        }
    }

    if success {
        let mut guard = chat_map.lock().unwrap();
        println!("Assigning partner {} to {}", partner, username);
        guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
    }
    success
}

// TODO: Receiving a message while in the middle of typing a message
//       inserts received message into middle of your message
fn chat(stream: TcpStream,
        username: String,
        partner: String,
        sender_to_router: Sender<Message>,
        receiver_from_router: Receiver<Message>) {

    // Send messages
    {
        let username = username.clone();
        let reader = BufReader::new(stream.try_clone().unwrap());
        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);

                let msg = Message::new(time::now().asctime().to_string(),
                                       username.to_string(),
                                       partner.to_string(),
                                       line.to_string());

                sender_to_router.send(msg).unwrap();
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        thread::spawn(move|| {
            loop {
                let stream = stream.try_clone().unwrap();
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}",
                            &username, msg.content());

                        receive_message(stream, msg);
                    },

                    Err(TryRecvError::Empty) => continue,

                    Err(TryRecvError::Disconnected) =>
                        panic!("User {} disconnected from router", &username)
                }
            }
        });
    }
}

fn receive_message(stream: TcpStream, msg: Message) {
    let mut stream = stream;
    let output = msg.from().to_string() + ": " + &msg.content()[..] + "\n";
    stream.write(&output.into_bytes()).unwrap();
}

#[cfg(test)]
mod client_tests {

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::channel;
    use chatmap::{ChatMap, ClientInfo};
    use super::try_select_partner;

    #[test]
    fn try_select_partner_test_success_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_success_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));
    }

    #[test]
    fn try_select_partner_test_succes_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: b can choose a when they are already partners
        //   (This allows someone to confirm a chat request)
        assert!(try_select_partner(&cm, "b".to_string(), "a".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: b already has a partner
        assert!(!try_select_partner(&cm, "c".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: Can't choose self as partner
        assert!(!try_select_partner(&cm, "c".to_string(), "c".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));

        // Fail: c already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "c".to_string()));

        // Fail: d already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "d".to_string()));
    }

    fn fixture() -> ChatMap {
        let mut cm = ChatMap::new();
        let (sender_to_a, _) = channel();
        cm.insert("a".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_a
        });
        let (sender_to_b, _) = channel();
        cm.insert("b".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_b
        });
        let (sender_to_c, _) = channel();
        cm.insert("c".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_c
        });
        let (sender_to_d, _) = channel();
        cm.insert("d".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_d
        });
        let (sender_to_e, _) = channel();
        cm.insert("e".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_e
        });
        cm
    }
}<|endoftext|>


<commit_before>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::*;
use command::{Command, parse_command};
use message::Message;

extern crate time;

pub fn create_client(stream: TcpStream,
                     sender_to_router: Sender<Message>,
                     chat_map: &Arc<Mutex<ChatMap>>) {

    println!("New client");
    let mut stream = stream;
    let welcome_msg = "Welcome to Smazy\nPlease enter a username:\n".to_string();
    stream.write(&welcome_msg.into_bytes()).expect("Error writing to stream");
    stream.flush().unwrap();

    let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
    let mut username = "".to_string();
    let (sender_to_client, receiver_from_router) = channel();

    loop {
        match reader.read_line(&mut username) {
            Ok(_) => {
                username = username.trim().to_string();

                {
                    let mut guard = chat_map.lock().expect("Error locking chatmap");

                    if is_valid_username(&*guard, username.to_string()) {

                        let client_info = ClientInfo{
                            partner: None,
                            sender_to_client: sender_to_client,
                            blocked_users: Vec::new(),
                        };

                        guard.insert(username.to_string(), client_info);

                        println!("New user: {}", username);
                        break;
                    }
                    else {
                        let invalid_msg = "Invalid username. Please try again.\n".to_string();
                        stream.write(&invalid_msg.into_bytes()).expect("Error writing to stream");
                        username = "".to_string();
                    }
                }
            },
            Err(e) => println!("{}", e)
        }
    }

    let partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
        username.to_string(), chat_map);

    chat(stream, username, partner, sender_to_router, receiver_from_router, chat_map);
}

fn display_available(stream: TcpStream, chat_map: &Arc<Mutex<ChatMap>>, username: String) {
    let mut stream = stream;

    let show_available_msg = "Here are the users available to chat:\n".to_string();
    stream.write(&show_available_msg.into_bytes()).expect("Error writing to stream");

    let avail;

    {
        let guard = chat_map.lock().expect("Error locking chatmap");
        avail = available_users(&*guard, username);
    }

    for user in avail {
        stream.write(&(user + "\n").to_string().into_bytes()).expect("Error writing to stream");
    }
}

fn choose_chat_partner(stream: TcpStream,
                       username: String,
                       chat_map: &Arc<Mutex<ChatMap>>) -> String {

    // Used to notify you if someone else starts chat with you
    // To begin, you have no partner
    let no_partner = Arc::new(Mutex::new(true));

    // Allow you to start a chat
    {
        let chat_map = chat_map.clone();
        let no_partner = no_partner.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");

        thread::spawn(move|| {

            loop {
                {
                    // Break if other thread has set no_partner to false
                    let guard = no_partner.lock().unwrap();
                    if !*guard {
                        break;
                    }
                }
                display_available(stream.try_clone().expect("Error cloning stream"),
                    &chat_map, username.to_string());

                let select_msg = "Select who you want to chat with:\n".to_string();
                stream.write(&select_msg.into_bytes()).expect("Error writing to stream");

                let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
                let mut partner = "".to_string();

                match reader.read_line(&mut partner) {
                    Ok(_) => {
                        partner = partner.trim().to_string();

                        if try_select_partner(&chat_map, username.to_string(), partner.to_string()) {
                            // If we successfully choose partner,
                            //   break out of this loop and end this thread.
                            //   Other thread will find that we have partner now
                            //   and will return that partner
                            break;
                        }
                    },
                    Err(e) => println!("{}", e)
                }
            }
        });
    }

    let handle;

    // Check if someone else has started a chat with you
    {
        let chat_map = chat_map.clone();
        let no_partner = no_partner.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");

        handle = thread::spawn(move|| {
            loop {
                let guard = chat_map.lock().expect("Error locking chatmap");

                if let &Some(ref p) = &guard.get(&username).unwrap().partner {
                    println!("{} chatting with {}", username, p);

                    let chatting_msg = "Now chatting with ".to_string()
                                        + p + ". Press enter to start chatting.\n";
                    stream.write(&chatting_msg.into_bytes()).expect("Error writing to stream");

                    {
                        let mut guard = no_partner.lock().unwrap();
                        *guard = false;
                    }
                    return p.to_string();
                }
            }
        });
    }

    // Return partner
    return handle.join().unwrap();
}

fn try_select_partner(chat_map: &Arc<Mutex<ChatMap>>,
                      username: String,
                      partner: String) -> bool {

    // Can't choose yourself as partner
    if &username[..] == &partner[..] {
        return false;
    }

    let mut guard = chat_map.lock().expect("Error locking chatmap");

    {
        // Not allowed to select someone you have blocked
        let my_client_info = guard.get(&username).unwrap();
        if my_client_info.blocked_users.contains(&partner) {
            return false;
        }
    }

    let success: bool;

    if let Some(clientinfo) = guard.get(&partner) {
        // If they haven't blocked you && they don't already have a partner
        if !clientinfo.blocked_users.contains(&username.to_string())
            && clientinfo.partner == None {
            success = true;
        } else {
            success = false;
        }
    } else {
        success = false
    }

    if success {
        println!("Assigning partners: {} and {}", username, partner);
        guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
        guard.get_mut(&partner).unwrap().partner = Some(username.to_string());
    }

    success
}

fn chat(stream: TcpStream,
        username: String,
        partner: String,
        sender_to_router: Sender<Message>,
        receiver_from_router: Receiver<Message>,
        chat_map: &Arc<Mutex<ChatMap>>) {

    let mut partner = partner.to_string();

    // Send messages
    {
        let username = username.clone();
        let chat_map = chat_map.clone();
        let stream = stream.try_clone().expect("Error cloning stream");
        let reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));

        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);

                let msg = Message::new(time::now().asctime().to_string(),
                                       username.to_string(),
                                       partner.to_string(),
                                       line.to_string());
                if line.len() > 0 {
                    if &line[0..1] == "/" {
                        match parse_command(msg.content().to_string()) {

                            Command::Quit => {
                                println!("Quit command");

                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                }

                                send_quit_message(username.to_string(), partner.to_string(), &sender_to_router);

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },

                            Command::DisplayAvailable => {
                                display_available(stream.try_clone().expect("Error cloning stream"),
                                    &chat_map, username.to_string());
                            },

                            Command::Block => {
                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                    let mut client_info = guard.get_mut(&username).unwrap();
                                    client_info.blocked_users.push(partner.to_string());
                                }

                                send_block_message(username.to_string(), partner.to_string(), &sender_to_router);

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },

                            Command::Unrecognized => println!("Unrecognized command")
                        }
                    } else {
                        sender_to_router.send(msg).expect("Error sending message");
                    }
                }
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        thread::spawn(move|| {
            loop {
                let stream = stream.try_clone().expect("Error cloning stream");
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}",
                            &username, msg.content());

                        receive_message(stream, msg);
                    },

                    Err(TryRecvError::Empty) => continue,

                    Err(TryRecvError::Disconnected) =>
                        panic!("User {} disconnected from router", username)
                }
            }
        });
    }
}

fn send_quit_message(username: String,
                     partner: String,
                     sender_to_router: &Sender<Message>) {

    let quit_msg = "I have quit out of the conversation. \
                    Type /q to quit.".to_string();

    let msg = Message::new(time::now().asctime().to_string(),
                           username, partner, quit_msg);

    sender_to_router.send(msg).expect("Error sending message");
}

fn send_block_message(username: String,
                      partner: String,
                      sender_to_router: &Sender<Message>) {

    let block_msg = "I have blocked you and quit out of the conversation. \
                     Type /q to quit.".to_string();

    let msg = Message::new(time::now().asctime().to_string(),
                           username, partner, block_msg);

    sender_to_router.send(msg).expect("Error sending message");
}

fn receive_message(stream: TcpStream, msg: Message) {
    let mut stream = stream;
    let output = msg.from().to_string() + ": " + &msg.content()[..] + "\n";
    stream.write(&output.into_bytes()).expect("Error writing to stream");
}

#[cfg(test)]
mod client_tests {

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::channel;
    use chatmap::{ChatMap, ClientInfo};

    use super::try_select_partner;

    #[test]
    fn try_select_partner_test_success_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_success_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: b already has a partner
        assert!(!try_select_partner(&cm, "c".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: a already has partner
        assert!(!try_select_partner(&cm, "b".to_string(), "a".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: Can't choose self as partner
        assert!(!try_select_partner(&cm, "c".to_string(), "c".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_4() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Fail: f not in chatmap
        assert!(!try_select_partner(&cm, "a".to_string(), "f".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_5() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));

        // Fail: c already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "c".to_string()));

        // Fail: d already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "d".to_string()));
    }

    fn fixture() -> ChatMap {
        let mut cm = ChatMap::new();
        let (sender_to_a, _) = channel();
        cm.insert("a".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_a,
            blocked_users: Vec::new()
        });
        let (sender_to_b, _) = channel();
        cm.insert("b".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_b,
            blocked_users: Vec::new()
        });
        let (sender_to_c, _) = channel();
        cm.insert("c".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_c,
            blocked_users: Vec::new()
        });
        let (sender_to_d, _) = channel();
        cm.insert("d".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_d,
            blocked_users: Vec::new()
        });
        let (sender_to_e, _) = channel();
        cm.insert("e".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_e,
            blocked_users: Vec::new()
        });
        cm
    }
}<commit_msg>Removed unnecessary scope<commit_after>use std::io::{BufRead, BufReader, Write};
use std::net::TcpStream;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{channel, Receiver, Sender, TryRecvError};
use std::thread;

use chatmap::{ChatMap, ClientInfo};
use command::Command;
use message::Message;

use chatmap::{available_users, is_valid_username, quit_conversation};
use command::parse_command;

extern crate time;

pub fn create_client(stream: TcpStream,
                     sender_to_router: Sender<Message>,
                     chat_map: &Arc<Mutex<ChatMap>>) {

    println!("New client");
    let mut stream = stream;
    let welcome_msg = "Welcome to Smazy\nPlease enter a username:\n".to_string();
    stream.write(&welcome_msg.into_bytes()).expect("Error writing to stream");
    stream.flush().unwrap();

    let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
    let mut username = "".to_string();
    let (sender_to_client, receiver_from_router) = channel();

    loop {
        match reader.read_line(&mut username) {
            Ok(_) => {
                username = username.trim().to_string();

                let mut guard = chat_map.lock().expect("Error locking chatmap");

                if is_valid_username(&*guard, username.to_string()) {

                    let client_info = ClientInfo{
                        partner: None,
                        sender_to_client: sender_to_client,
                        blocked_users: Vec::new(),
                    };

                    guard.insert(username.to_string(), client_info);

                    println!("New user: {}", username);
                    break;
                }
                else {
                    let invalid_msg = "Invalid username. Please try again.\n".to_string();
                    stream.write(&invalid_msg.into_bytes()).expect("Error writing to stream");
                    username = "".to_string();
                }
            },
            Err(e) => println!("{}", e)
        }
    }

    let partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
        username.to_string(), chat_map);

    chat(stream, username, partner, sender_to_router, receiver_from_router, chat_map);
}

fn display_available(stream: TcpStream, chat_map: &Arc<Mutex<ChatMap>>, username: String) {
    let mut stream = stream;

    let show_available_msg = "Here are the users available to chat:\n".to_string();
    stream.write(&show_available_msg.into_bytes()).expect("Error writing to stream");

    let avail;

    {
        let guard = chat_map.lock().expect("Error locking chatmap");
        avail = available_users(&*guard, username);
    }

    for user in avail {
        stream.write(&(user + "\n").to_string().into_bytes()).expect("Error writing to stream");
    }
}

fn choose_chat_partner(stream: TcpStream,
                       username: String,
                       chat_map: &Arc<Mutex<ChatMap>>) -> String {

    // Used to notify you if someone else starts chat with you
    // To begin, you have no partner
    let no_partner = Arc::new(Mutex::new(true));

    // Allow you to start a chat
    {
        let chat_map = chat_map.clone();
        let no_partner = no_partner.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");

        thread::spawn(move|| {

            loop {
                {
                    // Break if other thread has set no_partner to false
                    let guard = no_partner.lock().unwrap();
                    if !*guard {
                        break;
                    }
                }
                display_available(stream.try_clone().expect("Error cloning stream"),
                    &chat_map, username.to_string());

                let select_msg = "Select who you want to chat with:\n".to_string();
                stream.write(&select_msg.into_bytes()).expect("Error writing to stream");

                let mut reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));
                let mut partner = "".to_string();

                match reader.read_line(&mut partner) {
                    Ok(_) => {
                        partner = partner.trim().to_string();

                        if try_select_partner(&chat_map, username.to_string(), partner.to_string()) {
                            // If we successfully choose partner,
                            //   break out of this loop and end this thread.
                            //   Other thread will find that we have partner now
                            //   and will return that partner
                            break;
                        }
                    },
                    Err(e) => println!("{}", e)
                }
            }
        });
    }

    let handle;

    // Check if someone else has started a chat with you
    {
        let chat_map = chat_map.clone();
        let no_partner = no_partner.clone();
        let username = username.to_string();
        let mut stream = stream.try_clone().expect("Error cloning stream");

        handle = thread::spawn(move|| {
            loop {
                let guard = chat_map.lock().expect("Error locking chatmap");

                if let &Some(ref p) = &guard.get(&username).unwrap().partner {
                    println!("{} chatting with {}", username, p);

                    let chatting_msg = "Now chatting with ".to_string()
                                        + p + ". Press enter to start chatting.\n";
                    stream.write(&chatting_msg.into_bytes()).expect("Error writing to stream");

                    {
                        let mut guard = no_partner.lock().unwrap();
                        *guard = false;
                    }
                    return p.to_string();
                }
            }
        });
    }

    // Return partner
    return handle.join().unwrap();
}

fn try_select_partner(chat_map: &Arc<Mutex<ChatMap>>,
                      username: String,
                      partner: String) -> bool {

    // Can't choose yourself as partner
    if &username[..] == &partner[..] {
        return false;
    }

    let mut guard = chat_map.lock().expect("Error locking chatmap");

    {
        // Not allowed to select someone you have blocked
        let my_client_info = guard.get(&username).unwrap();
        if my_client_info.blocked_users.contains(&partner) {
            return false;
        }
    }

    let success: bool;

    if let Some(clientinfo) = guard.get(&partner) {
        // If they haven't blocked you && they don't already have a partner
        if !clientinfo.blocked_users.contains(&username.to_string())
            && clientinfo.partner == None {
            success = true;
        } else {
            success = false;
        }
    } else {
        success = false
    }

    if success {
        println!("Assigning partners: {} and {}", username, partner);
        guard.get_mut(&username).unwrap().partner = Some(partner.to_string());
        guard.get_mut(&partner).unwrap().partner = Some(username.to_string());
    }

    success
}

fn chat(stream: TcpStream,
        username: String,
        partner: String,
        sender_to_router: Sender<Message>,
        receiver_from_router: Receiver<Message>,
        chat_map: &Arc<Mutex<ChatMap>>) {

    let mut partner = partner.to_string();

    // Send messages
    {
        let username = username.clone();
        let chat_map = chat_map.clone();
        let stream = stream.try_clone().expect("Error cloning stream");
        let reader = BufReader::new(stream.try_clone().expect("Error cloning stream"));

        thread::spawn(move|| {
            let mut lines = reader.lines();
            while let Some(Ok(line)) = lines.next() {
                println!("{}",line);

                let msg = Message::new(time::now().asctime().to_string(),
                                       username.to_string(),
                                       partner.to_string(),
                                       line.to_string());
                if line.len() > 0 {
                    if &line[0..1] == "/" {
                        match parse_command(msg.content().to_string()) {

                            Command::Quit => {
                                println!("Quit command");

                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                }

                                send_quit_message(username.to_string(), partner.to_string(), &sender_to_router);

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },

                            Command::DisplayAvailable => {
                                display_available(stream.try_clone().expect("Error cloning stream"),
                                    &chat_map, username.to_string());
                            },

                            Command::Block => {
                                {
                                    let mut guard = chat_map.lock().expect("Error locking chatmap");
                                    quit_conversation(&mut *guard, username.to_string());
                                    let mut client_info = guard.get_mut(&username).unwrap();
                                    client_info.blocked_users.push(partner.to_string());
                                }

                                send_block_message(username.to_string(), partner.to_string(), &sender_to_router);

                                partner = choose_chat_partner(stream.try_clone().expect("Error cloning stream"),
                                    username.to_string(), &chat_map);
                            },

                            Command::Unrecognized => println!("Unrecognized command")
                        }
                    } else {
                        sender_to_router.send(msg).expect("Error sending message");
                    }
                }
            }
        });
    }

    // Receive Messages
    {
        let username = username.clone();
        thread::spawn(move|| {
            loop {
                let stream = stream.try_clone().expect("Error cloning stream");
                match receiver_from_router.try_recv() {
                    Ok(msg) => {
                        println!("User {} received message: {}",
                            &username, msg.content());

                        receive_message(stream, msg);
                    },

                    Err(TryRecvError::Empty) => continue,

                    Err(TryRecvError::Disconnected) =>
                        panic!("User {} disconnected from router", username)
                }
            }
        });
    }
}

fn send_quit_message(username: String,
                     partner: String,
                     sender_to_router: &Sender<Message>) {

    let quit_msg = "I have quit out of the conversation. \
                    Type /q to quit.".to_string();

    let msg = Message::new(time::now().asctime().to_string(),
                           username, partner, quit_msg);

    sender_to_router.send(msg).expect("Error sending message");
}

fn send_block_message(username: String,
                      partner: String,
                      sender_to_router: &Sender<Message>) {

    let block_msg = "I have blocked you and quit out of the conversation. \
                     Type /q to quit.".to_string();

    let msg = Message::new(time::now().asctime().to_string(),
                           username, partner, block_msg);

    sender_to_router.send(msg).expect("Error sending message");
}

fn receive_message(stream: TcpStream, msg: Message) {
    let mut stream = stream;
    let output = msg.from().to_string() + ": " + &msg.content()[..] + "\n";
    stream.write(&output.into_bytes()).expect("Error writing to stream");
}

#[cfg(test)]
mod client_tests {

    use std::sync::{Arc, Mutex};
    use std::sync::mpsc::channel;
    use chatmap::{ChatMap, ClientInfo};

    use super::try_select_partner;

    #[test]
    fn try_select_partner_test_success_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_success_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_1() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: b already has a partner
        assert!(!try_select_partner(&cm, "c".to_string(), "b".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_2() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: a already has partner
        assert!(!try_select_partner(&cm, "b".to_string(), "a".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_3() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Fail: Can't choose self as partner
        assert!(!try_select_partner(&cm, "c".to_string(), "c".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_4() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Fail: f not in chatmap
        assert!(!try_select_partner(&cm, "a".to_string(), "f".to_string()));
    }

    #[test]
    fn try_select_partner_test_fail_5() {
        let cm = Arc::new(Mutex::new(fixture()));

        // Success: a and b become partners
        assert!(try_select_partner(&cm, "a".to_string(), "b".to_string()));

        // Success: c and d become partners
        assert!(try_select_partner(&cm, "c".to_string(), "d".to_string()));

        // Fail: c already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "c".to_string()));

        // Fail: d already has a partner
        assert!(!try_select_partner(&cm, "e".to_string(), "d".to_string()));
    }

    fn fixture() -> ChatMap {
        let mut cm = ChatMap::new();
        let (sender_to_a, _) = channel();
        cm.insert("a".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_a,
            blocked_users: Vec::new()
        });
        let (sender_to_b, _) = channel();
        cm.insert("b".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_b,
            blocked_users: Vec::new()
        });
        let (sender_to_c, _) = channel();
        cm.insert("c".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_c,
            blocked_users: Vec::new()
        });
        let (sender_to_d, _) = channel();
        cm.insert("d".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_d,
            blocked_users: Vec::new()
        });
        let (sender_to_e, _) = channel();
        cm.insert("e".to_string(), ClientInfo{
            partner: None,
            sender_to_client: sender_to_e,
            blocked_users: Vec::new()
        });
        cm
    }
}<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};
use std::fs;
use clap::{App, Arg};
use ct_result::{CtResult, CtError, OkOr};

/// the configuration used to run `cpp-typecheck`
#[derive(Debug)]
pub struct Config {
    /// use this compiler for the type checking instead
    /// of the one specified in the database
    pub compiler: Option<String>,

    /// the C++ source file to type check
    pub cpp_file: PathBuf,

    /// the clang compilation databases to look up
    /// the C++ source file
    pub db_files: Vec<PathBuf>,

    /// forces the lookup in the database without considering
    /// the command cache
    pub no_cache: bool,

    /// forces the recaching of the command by doing the lookup
    /// in the database
    pub force_recache: bool
}

impl Config {
   pub fn from_command_args() -> CtResult<Config> {
       let matches = try!(App::new("cpp-typecheck")
           .about("Type check a C++ source file with a clang compilation database")
           .version(crate_version!())
           .author("Daniel Trstenjak <daniel.trstenjak@gmail.com>")
           .arg(Arg::with_name("compiler")
                .short("c")
                .long("compiler")
                .value_names(&["PATH"])
                .help("Use this compiler for the type checking instead of the one specified in the database")
                .takes_value(true))
           .arg(Arg::with_name("no-cache")
                .short("n")
                .long("no-cache")
                .help("Forces the lookup in the database without considering the command cache"))
           .arg(Arg::with_name("force-recache")
                .short("f")
                .long("force-recache")
                .help("Forces the recaching of the command by doing the lookup in the database"))
           .arg(Arg::with_name("SOURCE-FILE")
               .help("The C++ source file to type check")
               .required(true)
               .index(1))
           .arg(Arg::with_name("CLANG-DB")
               .help("The clang compilation database")
               .index(2)
               .multiple(true))
           .get_matches_safe());

       let mut cpp_file = PathBuf::from(try!(matches.value_of("SOURCE-FILE").ok_or("Missing C++ source file!")));
       try!(cpp_file.is_absolute().ok_or(format!("C++ source file '{}' has to have an absolute path!", cpp_file.display())));

       // if 'cpp_file' is a header file search for a C++ source file in the same directory
       {
           let header_file = cpp_file.clone();
           let ext = header_file.extension();
           let is_header_file = ext.is_none() || ext.unwrap() == "h";
           if is_header_file {
               let mut found_cpp_file: Option<PathBuf> = None;
               let cpp_exts = ["cpp", "cxx", "c"];
               for cpp_ext in &cpp_exts {
                   let file = header_file.with_extension(cpp_ext);
                   if file.is_file() {
                       found_cpp_file = Some(file);
                       break;
                   }
               }

               if found_cpp_file.is_none() {
                   return Err(CtError::from(format!("Couldn't find C++ source file for header '{}'!", header_file.display())));
               } else {
                   cpp_file = found_cpp_file.unwrap();
               }
           }
       }

       let db_files: Vec<PathBuf> = {
           if let Some(values) = matches.values_of("CLANG-DB") {
               values.map(PathBuf::from).collect()
           } else {
               let dir = try!(cpp_file.parent()
                  .ok_or(format!("Couldn't get directory of source file '{}'!", cpp_file.display())));

               vec![try!(find_db(&dir))]
           }
       };

       try!((! db_files.is_empty()).ok_or("Missing clang compilation database!"));

       Ok(Config {
           compiler: matches.value_of("compiler").map(String::from),
           cpp_file: cpp_file,
           db_files: db_files,
           no_cache: matches.is_present("no-cache"),
           force_recache: matches.is_present("force-recache")
       })
   }
}

/// Searches for a `compile_commands.json` file starting at `start_dir` and continuing the search
/// upwards the directory tree until the file is found.
fn find_db(start_dir: &Path) -> CtResult<PathBuf> {
    let mut dir = start_dir.to_path_buf();
    loop {
        if let Ok(files) = fs::read_dir(&dir) {
            for file in files {
                if let Ok(file) = file {
                    let path = file.path();
                    if path.is_file() {
                        if let Some("compile_commands.json") = path.file_name().and_then(|s| s.to_str()) {
                            return Ok(path);
                        }
                    }
                }
            }
        }

        try!(dir.pop().ok_or(format!("Couldn't find 'compile_commands.json' starting at directory '{}'!",
                                     start_dir.display())));
    }
}
<commit_msg>Support headers with 'hpp' extension<commit_after>use std::path::{Path, PathBuf};
use std::fs;
use clap::{App, Arg};
use ct_result::{CtResult, CtError, OkOr};

/// the configuration used to run `cpp-typecheck`
#[derive(Debug)]
pub struct Config {
    /// use this compiler for the type checking instead
    /// of the one specified in the database
    pub compiler: Option<String>,

    /// the C++ source file to type check
    pub cpp_file: PathBuf,

    /// the clang compilation databases to look up
    /// the C++ source file
    pub db_files: Vec<PathBuf>,

    /// forces the lookup in the database without considering
    /// the command cache
    pub no_cache: bool,

    /// forces the recaching of the command by doing the lookup
    /// in the database
    pub force_recache: bool
}

impl Config {
   pub fn from_command_args() -> CtResult<Config> {
       let matches = try!(App::new("cpp-typecheck")
           .about("Type check a C++ source file with a clang compilation database")
           .version(crate_version!())
           .author("Daniel Trstenjak <daniel.trstenjak@gmail.com>")
           .arg(Arg::with_name("compiler")
                .short("c")
                .long("compiler")
                .value_names(&["PATH"])
                .help("Use this compiler for the type checking instead of the one specified in the database")
                .takes_value(true))
           .arg(Arg::with_name("no-cache")
                .short("n")
                .long("no-cache")
                .help("Forces the lookup in the database without considering the command cache"))
           .arg(Arg::with_name("force-recache")
                .short("f")
                .long("force-recache")
                .help("Forces the recaching of the command by doing the lookup in the database"))
           .arg(Arg::with_name("SOURCE-FILE")
               .help("The C++ source file to type check")
               .required(true)
               .index(1))
           .arg(Arg::with_name("CLANG-DB")
               .help("The clang compilation database")
               .index(2)
               .multiple(true))
           .get_matches_safe());

       let mut cpp_file = PathBuf::from(try!(matches.value_of("SOURCE-FILE").ok_or("Missing C++ source file!")));
       try!(cpp_file.is_absolute().ok_or(format!("C++ source file '{}' has to have an absolute path!", cpp_file.display())));

       // if 'cpp_file' is a header file search for a C++ source file in the same directory
       {
           let header_file = cpp_file.clone();
           let is_header_file = if let Some(ext) = header_file.extension() {
               ext == "h" || ext == "hpp"
           } else {
               // no extension, assume a header file
               true
           };

           if is_header_file {
               let mut found_cpp_file: Option<PathBuf> = None;
               let cpp_exts = ["cpp", "cxx", "c"];
               for cpp_ext in &cpp_exts {
                   let file = header_file.with_extension(cpp_ext);
                   if file.is_file() {
                       found_cpp_file = Some(file);
                       break;
                   }
               }

               if found_cpp_file.is_none() {
                   return Err(CtError::from(format!("Couldn't find C++ source file for header '{}'!", header_file.display())));
               } else {
                   cpp_file = found_cpp_file.unwrap();
               }
           }
       }

       let db_files: Vec<PathBuf> = {
           if let Some(values) = matches.values_of("CLANG-DB") {
               values.map(PathBuf::from).collect()
           } else {
               let dir = try!(cpp_file.parent()
                  .ok_or(format!("Couldn't get directory of source file '{}'!", cpp_file.display())));

               vec![try!(find_db(&dir))]
           }
       };

       try!((! db_files.is_empty()).ok_or("Missing clang compilation database!"));

       Ok(Config {
           compiler: matches.value_of("compiler").map(String::from),
           cpp_file: cpp_file,
           db_files: db_files,
           no_cache: matches.is_present("no-cache"),
           force_recache: matches.is_present("force-recache")
       })
   }
}

/// Searches for a `compile_commands.json` file starting at `start_dir` and continuing the search
/// upwards the directory tree until the file is found.
fn find_db(start_dir: &Path) -> CtResult<PathBuf> {
    let mut dir = start_dir.to_path_buf();
    loop {
        if let Ok(files) = fs::read_dir(&dir) {
            for file in files {
                if let Ok(file) = file {
                    let path = file.path();
                    if path.is_file() {
                        if let Some("compile_commands.json") = path.file_name().and_then(|s| s.to_str()) {
                            return Ok(path);
                        }
                    }
                }
            }
        }

        try!(dir.pop().ok_or(format!("Couldn't find 'compile_commands.json' starting at directory '{}'!",
                                     start_dir.display())));
    }
}
<|endoftext|>


<commit_before>use anyhow::{Context, Result};
use serde_derive::Deserialize;
use std::path::Path;

static GALLERY_FOLDER_NAME: &str = "images";
static OUTPUT_FOLDER_NAME: &str = "_output";

#[derive(Debug, Deserialize, Default)]
pub struct Config {
    pub title: Option<String>,
    pub template_dir: Option<String>,
    pub output_dir: Option<String>,
    pub logging: Option<LogKind>,
    pub remove_numbered_prefix: Option<bool>,
    pub copy_dirs: Option<Vec<String>>,
    pub gallery: Option<Gallery>,
    pub sync_settings: Option<SyncSettings>,
}

#[derive(Debug, Deserialize)]
pub struct Gallery {
    pub img_dir: Option<String>,
    pub img_width: u32,
    pub img_height: u32,
    pub thumb_width: u32,
    pub thumb_height: u32,
}

#[derive(PartialEq, Eq, Clone, Copy, Debug, Deserialize)]
pub enum FtpProtocol {
    Ftp,
    Sftp,
}

#[derive(PartialEq, Eq, Clone, Copy, Debug, Deserialize)]
pub enum LogKind {
    Stdout,
    File,
}

#[derive(Debug, Deserialize)]
pub struct SyncSettings {
    pub ftp_server: String,
    pub ftp_port: Option<u32>,
    pub ftp_protocol: Option<FtpProtocol>,
    pub ftp_user: String,
    pub ftp_target_dir: Option<String>,
    pub ftp_overwrite: Option<bool>,
}

impl Config {
    pub fn load(path: &Path) -> Result<Config> {
        let mut input = String::new();
        std::fs::File::open(path.join("config.toml").as_path())
            .and_then(|mut f| std::io::Read::read_to_string(&mut f, &mut input))
            .context("couldn't find or read file 'config.toml'")?;
        let mut conf =
            toml::from_str::<Config>(input.as_str()).context("parsing 'config.toml' failed")?;
        conf.resolve_paths(path);
        Ok(conf)
    }

    pub fn resolve_paths(&mut self, base_path: &Path) {
        let output_path = base_path.join(&self.output_dir.as_deref().unwrap_or(OUTPUT_FOLDER_NAME));
        self.output_dir = Some(
            output_path
                .as_path()
                .to_str()
                .expect("Could not resolve path to template directory")
                .to_owned(),
        );
        if self.gallery.is_none() {
            self.gallery = Some(Gallery {
                img_dir: Some(GALLERY_FOLDER_NAME.to_string()),
                img_width: 600,
                img_height: 800,
                thumb_width: 90,
                thumb_height: 90,
            })
        } else if self.gallery.as_ref().unwrap().img_dir.is_none() {
            self.gallery.as_mut().unwrap().img_dir = Some(GALLERY_FOLDER_NAME.to_string());
        }
        if self.template_dir.is_some() {
            let template_path = base_path.join(&self.template_dir.as_ref().unwrap().as_str());
            assert!(
                template_path.exists(),
                "The template directory '{}' does not exist in your project '{}'",
                &self.template_dir.as_ref().unwrap().as_str(),
                base_path.display()
            );
            self.template_dir = Some(
                template_path
                    .as_path()
                    .to_str()
                    .expect("Could not resolve path to template directory")
                    .to_owned(),
            );
        }
        if let Some(ref copy_dirs) = self.copy_dirs {
            let mut new_copy_dirs = Vec::<String>::new();
            for copy_dir in copy_dirs {
                let path = base_path.join(copy_dir.as_str());
                assert!(
                    path.exists(),
                    "The directory '{}' does not exist in your project '{}'",
                    copy_dir.as_str(),
                    base_path.display()
                );
                new_copy_dirs.push(
                    path.as_path()
                        .to_str()
                        .expect("Could not resolve path to copy directory")
                        .to_owned(),
                );
            }
            self.copy_dirs = Some(new_copy_dirs);
        }
        if let Some(ref mut sync_settings) = self.sync_settings {
            if let Some(ref mut ftp_target_dir) = sync_settings.ftp_target_dir {
                if !ftp_target_dir.starts_with('/') {
                    println!("ftp_target_dir {ftp_target_dir} is treated as absolute path!");
                    ftp_target_dir.insert(0, '/');
                }
            }
        }
    }

    pub fn print(&self) {
        use term_painter::Attr::Bold;
        use term_painter::ToStyle;
        let none_string = "None".to_string();
        println!("Title : {}", self.title.as_ref().unwrap_or(&none_string));
        println!(
            "Template directory : {}",
            self.template_dir.as_ref().unwrap_or(&none_string)
        );
        println!(
            "Output directory : {}",
            self.output_dir
                .as_ref()
                .unwrap_or(&OUTPUT_FOLDER_NAME.to_string(),)
        );
        println!("{}", Bold.paint("SyncSettings"));
        if self.sync_settings.is_some() {
            let sync_settings = self.sync_settings.as_ref().unwrap();
            println!("  FTP server: {}", sync_settings.ftp_server);
            println!("  FTP port: {}", sync_settings.ftp_port.unwrap_or(21));
            println!("  FTP user: {}", sync_settings.ftp_user);
            println!(
                "  FTP overwrite: {}",
                sync_settings.ftp_overwrite.unwrap_or(false)
            );
        }
        if self.gallery.is_some() {
            let gallery = self.gallery.as_ref().unwrap();
            println!("{}", Bold.paint("Gallery"));
            println!(
                "  image directory: {}",
                gallery.img_dir.as_ref().unwrap_or(&"images".to_string())
            );
            println!(
                "  image size : {} x {}",
                gallery.img_width, gallery.img_height
            );
            println!(
                "  thumb size : {} x {}",
                gallery.thumb_width, gallery.thumb_height
            );
        }
    }
}
<commit_msg>improve error message<commit_after>use anyhow::{Context, Result};
use serde_derive::Deserialize;
use std::path::Path;

static GALLERY_FOLDER_NAME: &str = "images";
static OUTPUT_FOLDER_NAME: &str = "_output";

#[derive(Debug, Deserialize, Default)]
pub struct Config {
    pub title: Option<String>,
    pub template_dir: Option<String>,
    pub output_dir: Option<String>,
    pub logging: Option<LogKind>,
    pub remove_numbered_prefix: Option<bool>,
    pub copy_dirs: Option<Vec<String>>,
    pub gallery: Option<Gallery>,
    pub sync_settings: Option<SyncSettings>,
}

#[derive(Debug, Deserialize)]
pub struct Gallery {
    pub img_dir: Option<String>,
    pub img_width: u32,
    pub img_height: u32,
    pub thumb_width: u32,
    pub thumb_height: u32,
}

#[derive(PartialEq, Eq, Clone, Copy, Debug, Deserialize)]
pub enum FtpProtocol {
    Ftp,
    Sftp,
}

#[derive(PartialEq, Eq, Clone, Copy, Debug, Deserialize)]
pub enum LogKind {
    Stdout,
    File,
}

#[derive(Debug, Deserialize)]
pub struct SyncSettings {
    pub ftp_server: String,
    pub ftp_port: Option<u32>,
    pub ftp_protocol: Option<FtpProtocol>,
    pub ftp_user: String,
    pub ftp_target_dir: Option<String>,
    pub ftp_overwrite: Option<bool>,
}

impl Config {
    pub fn load(path: &Path) -> Result<Config> {
        let mut input = String::new();
        std::fs::File::open(path.join("config.toml").as_path())
            .and_then(|mut f| std::io::Read::read_to_string(&mut f, &mut input))
            .context("couldn't find or read file 'config.toml'")?;
        let mut conf =
            toml::from_str::<Config>(input.as_str()).context("parsing 'config.toml' failed")?;
        conf.resolve_paths(path);
        Ok(conf)
    }

    pub fn resolve_paths(&mut self, base_path: &Path) {
        let output_path = base_path.join(&self.output_dir.as_deref().unwrap_or(OUTPUT_FOLDER_NAME));
        self.output_dir = Some(
            output_path
                .as_path()
                .to_str()
                .expect("Could not resolve path to template directory")
                .to_owned(),
        );
        if self.gallery.is_none() {
            self.gallery = Some(Gallery {
                img_dir: Some(GALLERY_FOLDER_NAME.to_string()),
                img_width: 600,
                img_height: 800,
                thumb_width: 90,
                thumb_height: 90,
            })
        } else if self.gallery.as_ref().unwrap().img_dir.is_none() {
            self.gallery.as_mut().unwrap().img_dir = Some(GALLERY_FOLDER_NAME.to_string());
        }
        if self.template_dir.is_some() {
            let template_path = base_path.join(&self.template_dir.as_ref().unwrap().as_str());
            assert!(
                template_path.exists(),
                "The template directory '{}' does not exist in your project '{}'",
                &self.template_dir.as_ref().unwrap().as_str(),
                base_path.display()
            );
            self.template_dir = Some(
                template_path
                    .as_path()
                    .to_str()
                    .expect("Could not resolve path to template directory")
                    .to_owned(),
            );
        }
        if let Some(ref copy_dirs) = self.copy_dirs {
            let mut new_copy_dirs = Vec::<String>::new();
            for copy_dir in copy_dirs {
                let path = base_path.join(copy_dir.as_str());
                assert!(
                    path.exists(),
                    "The directory '{}' does not exist in your project '{}'",
                    copy_dir.as_str(),
                    base_path.display()
                );
                new_copy_dirs.push(
                    path.as_path()
                        .to_str()
                        .expect("Could not resolve path to copy directory")
                        .to_owned(),
                );
            }
            self.copy_dirs = Some(new_copy_dirs);
        }
        if let Some(ref mut sync_settings) = self.sync_settings {
            if let Some(ref mut ftp_target_dir) = sync_settings.ftp_target_dir {
                if !ftp_target_dir.starts_with('/') {
                    println!("ftp_target_dir contains no absolute path '/' will be prepended to '{ftp_target_dir}'");
                    ftp_target_dir.insert(0, '/');
                }
            }
        }
    }

    pub fn print(&self) {
        use term_painter::Attr::Bold;
        use term_painter::ToStyle;
        let none_string = "None".to_string();
        println!("Title : {}", self.title.as_ref().unwrap_or(&none_string));
        println!(
            "Template directory : {}",
            self.template_dir.as_ref().unwrap_or(&none_string)
        );
        println!(
            "Output directory : {}",
            self.output_dir
                .as_ref()
                .unwrap_or(&OUTPUT_FOLDER_NAME.to_string(),)
        );
        println!("{}", Bold.paint("SyncSettings"));
        if self.sync_settings.is_some() {
            let sync_settings = self.sync_settings.as_ref().unwrap();
            println!("  FTP server: {}", sync_settings.ftp_server);
            println!("  FTP port: {}", sync_settings.ftp_port.unwrap_or(21));
            println!("  FTP user: {}", sync_settings.ftp_user);
            println!(
                "  FTP overwrite: {}",
                sync_settings.ftp_overwrite.unwrap_or(false)
            );
        }
        if self.gallery.is_some() {
            let gallery = self.gallery.as_ref().unwrap();
            println!("{}", Bold.paint("Gallery"));
            println!(
                "  image directory: {}",
                gallery.img_dir.as_ref().unwrap_or(&"images".to_string())
            );
            println!(
                "  image size : {} x {}",
                gallery.img_width, gallery.img_height
            );
            println!(
                "  thumb size : {} x {}",
                gallery.thumb_width, gallery.thumb_height
            );
        }
    }
}
<|endoftext|>


<commit_before>//! Configuration module for Ataxia
use serde_derive::Deserialize;
use std::fs::File;
use std::io::Read;
use std::path::Path;

/// Config structure for holding internal and external configuration data
#[derive(Deserialize, Debug)]
pub struct Config {
    #[serde(default)]
    http_addr: String,
    #[serde(default)]
    telnet_addr: String,
    #[serde(default)]
    internal_addr: String,
    pid_file: String,
    log_file: String,
    #[serde(default)]
    debug: bool,
    #[serde(default)]
    verbose: bool,
}

impl Config {
    #![allow(clippy::new_ret_no_self)]
    /// Returns a new Config
    /// Read configuration from the file path specified in the argument structure.
    ///
    /// # Arguments
    ///
    /// * `matches` - A clap::ArgMatches structure containing command-line arguments and default values
    ///
    /// # Errors
    ///
    /// * Returns std::io::Error if the config file can't be opened or read
    /// * Returns toml::de::Error if TOML parsing fails
    ///
    pub fn new(matches: &clap::ArgMatches<'_>) -> Result<Self, failure::Error> {
        let path = Path::new(matches.value_of("config").unwrap_or("data/ataxia.toml"));

        let mut input = String::new();
        File::open(path)?.read_to_string(&mut input)?;
        let mut config = toml::from_str::<Self>(&input)?;

        if let Some(pid_file) = matches.value_of("pid_file") {
            config.set_pid_file(pid_file);
        }

        if let Some(http_addr) = matches.value_of("http_addr") {
            config.set_http_addr(http_addr);
        }

        if let Some(telnet_addr) = matches.value_of("telnet_addr") {
            config.set_telnet_addr(telnet_addr);
        }

        if let Some(internal_addr) = matches.value_of("internal_addr") {
            config.set_internal_addr(internal_addr);
        }

        config.debug = match matches.occurrences_of("debug") {
            0 => false,
            1 | _ => true,
        };

        config.verbose = match matches.occurrences_of("verbose") {
            0 => false,
            1 | _ => true,
        };

        Ok(config)
    }

    /// Returns the listen address of the network proxy process for http connections
    pub fn http_addr(&self) -> &str {
        self.http_addr.as_ref()
    }
    /// Set the listen address of the network proxy process for http connections
    pub fn set_http_addr(&mut self, addr: &str) {
        self.http_addr = addr.to_string();
    }

    /// Returns the listen address of the network proxy process for telnet connections
    pub fn telnet_addr(&self) -> &str {
        self.telnet_addr.as_ref()
    }
    /// Set the listen address of the network proxy process for telnet connections
    pub fn set_telnet_addr(&mut self, addr: &str) {
        self.telnet_addr = addr.to_string();
    }

    /// Returns the listen address of the network proxy process for internal connections
    pub fn internal_addr(&self) -> &str {
        self.internal_addr.as_ref()
    }
    /// Set the listen address of the network proxy process for internal connections
    pub fn set_internal_addr(&mut self, addr: &str) {
        self.internal_addr = addr.to_string();
    }

    /// Returns the file path to the pid file
    pub fn pid_file(&self) -> &str {
        self.pid_file.as_ref()
    }
    /// Set the file path to the pid file
    pub fn set_pid_file(&mut self, file: &str) {
        self.pid_file = file.to_string();
    }

    /// Returns the file path to the log file
    pub fn log_file(&self) -> &str {
        self.log_file.as_ref()
    }

    /// Returns true if the debug CLI flag was specified
    pub fn debug(&self) -> bool {
        self.debug
    }

    /// Returns true if the verbose CLI flag was specified
    pub fn verbose(&self) -> bool {
        self.verbose
    }
}
<commit_msg>Minor cleanup of the config module.<commit_after>//! Configuration module for Ataxia
use serde::Deserialize;
use std::fs::File;
use std::io::Read;
use std::path::Path;

/// Config structure for holding internal and external configuration data
#[derive(Deserialize, Debug)]
pub struct Config {
    #[serde(default)]
    http_addr: String,
    #[serde(default)]
    telnet_addr: String,
    #[serde(default)]
    internal_addr: String,
    pid_file: String,
    log_file: String,
    #[serde(default)]
    debug: bool,
    #[serde(default)]
    verbose: bool,
}

impl Config {
    #![allow(clippy::new_ret_no_self)]
    /// Returns a new Config
    /// Read configuration from the file path specified in the Clap arguments struct.
    ///
    /// # Arguments
    ///
    /// * `matches` - A clap::ArgMatches structure containing command-line arguments and default values
    ///
    /// # Errors
    ///
    /// * Returns std::io::Error if the config file can't be opened or read
    /// * Returns toml::de::Error if TOML parsing fails
    ///
    pub fn new(matches: &clap::ArgMatches<'_>) -> Result<Self, failure::Error> {
        let path = Path::new(matches.value_of("config").unwrap_or("data/ataxia.toml"));

        let mut input = String::new();
        File::open(path)?.read_to_string(&mut input)?;
        let mut config = toml::from_str::<Self>(&input)?;

        if let Some(pid_file) = matches.value_of("pid_file") {
            config.pid_file = pid_file.to_string();
        }

        if let Some(http_addr) = matches.value_of("http_addr") {
            config.http_addr = http_addr.to_string();
        }

        if let Some(telnet_addr) = matches.value_of("telnet_addr") {
            config.telnet_addr = telnet_addr.to_string();
        }

        if let Some(internal_addr) = matches.value_of("internal_addr") {
            config.internal_addr = internal_addr.to_string();
        }

        config.debug = match matches.occurrences_of("debug") {
            0 => false,
            1 | _ => true,
        };

        config.verbose = match matches.occurrences_of("verbose") {
            0 => false,
            1 | _ => true,
        };

        Ok(config)
    }

    /// Returns the listen address of the network proxy process for http connections
    pub fn http_addr(&self) -> &str {
        self.http_addr.as_ref()
    }
    /// Set the listen address of the network proxy process for http connections
    pub fn set_http_addr(&mut self, addr: String) {
        self.http_addr = addr;
    }

    /// Returns the listen address of the network proxy process for telnet connections
    pub fn telnet_addr(&self) -> &str {
        self.telnet_addr.as_ref()
    }
    /// Set the listen address of the network proxy process for telnet connections
    pub fn set_telnet_addr(&mut self, addr: String) {
        self.telnet_addr = addr;
    }

    /// Returns the listen address of the network proxy process for internal connections
    pub fn internal_addr(&self) -> &str {
        self.internal_addr.as_ref()
    }
    /// Set the listen address of the network proxy process for internal connections
    pub fn set_internal_addr(&mut self, addr: String) {
        self.internal_addr = addr;
    }

    /// Returns the file path to the pid file
    pub fn pid_file(&self) -> &str {
        self.pid_file.as_ref()
    }
    /// Set the file path to the pid file
    pub fn set_pid_file(&mut self, file: String) {
        self.pid_file = file;
    }

    /// Returns the file path to the log file
    pub fn log_file(&self) -> &str {
        self.log_file.as_ref()
    }

    /// Returns true if the debug CLI flag was specified
    pub fn debug(&self) -> bool {
        self.debug
    }

    /// Returns true if the verbose CLI flag was specified
    pub fn verbose(&self) -> bool {
        self.verbose
    }
}
<|endoftext|>


<commit_before>use serde::de::{Deserialize, Deserializer};
use serde::ser::{Serialize, Serializer};
use std::collections::HashMap;
use std::fmt::Debug;
use std::ops::Deref;
use std::str::FromStr;

use error::*;
use ser::ConfigSerializer;
use source::Source;

use path;
use value::{Table, Value, ValueKind};

#[derive(Clone, Debug)]
enum ConfigKind {
    // A mutable configuration. This is the default.
    Mutable {
        defaults: HashMap<path::Expression, Value>,
        overrides: HashMap<path::Expression, Value>,
        sources: Vec<Box<dyn Source + Send + Sync>>,
    },

    // A frozen configuration.
    // Configuration can no longer be mutated.
    Frozen,
}

impl Default for ConfigKind {
    fn default() -> Self {
        ConfigKind::Mutable {
            defaults: HashMap::new(),
            overrides: HashMap::new(),
            sources: Vec::new(),
        }
    }
}

/// A prioritized configuration repository. It maintains a set of
/// configuration sources, fetches values to populate those, and provides
/// them according to the source's priority.
#[derive(Default, Clone, Debug)]
pub struct Config {
    kind: ConfigKind,

    /// Root of the cached configuration.
    pub cache: Value,
}

impl Config {
    pub fn new() -> Self {
        Self {
            kind: ConfigKind::default(),
            // Config root should be instantiated as an empty table
            // to avoid deserialization errors.
            cache: Value::new(None, Table::new()),
        }
    }

    /// Merge in a configuration property source.
    pub fn merge<T>(&mut self, source: T) -> Result<&mut Config>
    where
        T: 'static,
        T: Source + Send + Sync,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut sources, ..
            } => {
                sources.push(Box::new(source));
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        }

        self.refresh()
    }

    /// Merge in a configuration property source.
    pub fn with_merged<T>(mut self, source: T) -> Result<Self>
    where
        T: 'static,
        T: Source + Send + Sync,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut sources, ..
            } => {
                sources.push(Box::new(source));
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        }

        self.refresh()?;
        Ok(self)
    }

    /// Refresh the configuration cache with fresh
    /// data from added sources.
    ///
    /// Configuration is automatically refreshed after a mutation
    /// operation (`set`, `merge`, `set_default`, etc.).
    pub fn refresh(&mut self) -> Result<&mut Config> {
        self.cache = match self.kind {
            // TODO: We need to actually merge in all the stuff
            ConfigKind::Mutable {
                ref overrides,
                ref sources,
                ref defaults,
            } => {
                let mut cache: Value = HashMap::<String, Value>::new().into();

                // Add defaults
                for (key, val) in defaults {
                    key.set(&mut cache, val.clone());
                }

                // Add sources
                sources.collect_to(&mut cache)?;

                // Add overrides
                for (key, val) in overrides {
                    key.set(&mut cache, val.clone());
                }

                cache
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        };

        Ok(self)
    }

    pub fn set_default<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut defaults, ..
            } => {
                defaults.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    /// Set an overwrite
    ///
    /// This function sets an overwrite value.
    /// The overwrite `value` is written to the `key` location on every `refresh()`
    ///
    /// # Warning
    ///
    /// Errors if config is frozen
    pub fn set<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut overrides, ..
            } => {
                overrides.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    pub fn set_once(&mut self, key: &str, value: Value) -> Result<()> {
        let expr: path::Expression = key.parse()?;

        // Traverse the cache using the path to (possibly) retrieve a value
        if let Some(ref mut val) = expr.get_mut(&mut self.cache) {
            **val = value;
        } else {
            expr.set(&mut self.cache, value);
        }
        Ok(())
    }

    pub fn get<'de, T: Deserialize<'de>>(&self, key: &str) -> Result<T> {
        // Parse the key into a path expression
        let expr: path::Expression = key.parse()?;

        // Traverse the cache using the path to (possibly) retrieve a value
        let value = expr.get(&self.cache).cloned();

        match value {
            Some(value) => {
                // Deserialize the received value into the requested type
                T::deserialize(value).map_err(|e| e.extend_with_key(key))
            }

            None => Err(ConfigError::NotFound(key.into())),
        }
    }

    pub fn get_str(&self, key: &str) -> Result<String> {
        self.get(key).and_then(Value::into_str)
    }

    pub fn get_int(&self, key: &str) -> Result<i64> {
        self.get(key).and_then(Value::into_int)
    }

    pub fn get_float(&self, key: &str) -> Result<f64> {
        self.get(key).and_then(Value::into_float)
    }

    pub fn get_bool(&self, key: &str) -> Result<bool> {
        self.get(key).and_then(Value::into_bool)
    }

    pub fn get_table(&self, key: &str) -> Result<HashMap<String, Value>> {
        self.get(key).and_then(Value::into_table)
    }

    pub fn get_array(&self, key: &str) -> Result<Vec<Value>> {
        self.get(key).and_then(Value::into_array)
    }

    /// Attempt to deserialize the entire configuration into the requested type.
    pub fn try_into<'de, T: Deserialize<'de>>(self) -> Result<T> {
        T::deserialize(self)
    }

    /// Attempt to serialize the entire configuration from the given type.
    pub fn try_from<T: Serialize>(from: &T) -> Result<Self> {
        let mut serializer = ConfigSerializer::default();
        from.serialize(&mut serializer)?;
        Ok(serializer.output)
    }

    #[deprecated(since = "0.7.0", note = "please use 'try_into' instead")]
    pub fn deserialize<'de, T: Deserialize<'de>>(self) -> Result<T> {
        self.try_into()
    }
}

impl Source for Config {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        self.cache.clone().into_table()
    }
}
<commit_msg>Add doc: Config::set_default()<commit_after>use serde::de::{Deserialize, Deserializer};
use serde::ser::{Serialize, Serializer};
use std::collections::HashMap;
use std::fmt::Debug;
use std::ops::Deref;
use std::str::FromStr;

use error::*;
use ser::ConfigSerializer;
use source::Source;

use path;
use value::{Table, Value, ValueKind};

#[derive(Clone, Debug)]
enum ConfigKind {
    // A mutable configuration. This is the default.
    Mutable {
        defaults: HashMap<path::Expression, Value>,
        overrides: HashMap<path::Expression, Value>,
        sources: Vec<Box<dyn Source + Send + Sync>>,
    },

    // A frozen configuration.
    // Configuration can no longer be mutated.
    Frozen,
}

impl Default for ConfigKind {
    fn default() -> Self {
        ConfigKind::Mutable {
            defaults: HashMap::new(),
            overrides: HashMap::new(),
            sources: Vec::new(),
        }
    }
}

/// A prioritized configuration repository. It maintains a set of
/// configuration sources, fetches values to populate those, and provides
/// them according to the source's priority.
#[derive(Default, Clone, Debug)]
pub struct Config {
    kind: ConfigKind,

    /// Root of the cached configuration.
    pub cache: Value,
}

impl Config {
    pub fn new() -> Self {
        Self {
            kind: ConfigKind::default(),
            // Config root should be instantiated as an empty table
            // to avoid deserialization errors.
            cache: Value::new(None, Table::new()),
        }
    }

    /// Merge in a configuration property source.
    pub fn merge<T>(&mut self, source: T) -> Result<&mut Config>
    where
        T: 'static,
        T: Source + Send + Sync,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut sources, ..
            } => {
                sources.push(Box::new(source));
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        }

        self.refresh()
    }

    /// Merge in a configuration property source.
    pub fn with_merged<T>(mut self, source: T) -> Result<Self>
    where
        T: 'static,
        T: Source + Send + Sync,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut sources, ..
            } => {
                sources.push(Box::new(source));
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        }

        self.refresh()?;
        Ok(self)
    }

    /// Refresh the configuration cache with fresh
    /// data from added sources.
    ///
    /// Configuration is automatically refreshed after a mutation
    /// operation (`set`, `merge`, `set_default`, etc.).
    pub fn refresh(&mut self) -> Result<&mut Config> {
        self.cache = match self.kind {
            // TODO: We need to actually merge in all the stuff
            ConfigKind::Mutable {
                ref overrides,
                ref sources,
                ref defaults,
            } => {
                let mut cache: Value = HashMap::<String, Value>::new().into();

                // Add defaults
                for (key, val) in defaults {
                    key.set(&mut cache, val.clone());
                }

                // Add sources
                sources.collect_to(&mut cache)?;

                // Add overrides
                for (key, val) in overrides {
                    key.set(&mut cache, val.clone());
                }

                cache
            }

            ConfigKind::Frozen => {
                return Err(ConfigError::Frozen);
            }
        };

        Ok(self)
    }

    /// Set a default `value` at `key`
    pub fn set_default<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut defaults, ..
            } => {
                defaults.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    /// Set an overwrite
    ///
    /// This function sets an overwrite value.
    /// The overwrite `value` is written to the `key` location on every `refresh()`
    ///
    /// # Warning
    ///
    /// Errors if config is frozen
    pub fn set<T>(&mut self, key: &str, value: T) -> Result<&mut Config>
    where
        T: Into<Value>,
    {
        match self.kind {
            ConfigKind::Mutable {
                ref mut overrides, ..
            } => {
                overrides.insert(key.parse()?, value.into());
            }

            ConfigKind::Frozen => return Err(ConfigError::Frozen),
        };

        self.refresh()
    }

    pub fn set_once(&mut self, key: &str, value: Value) -> Result<()> {
        let expr: path::Expression = key.parse()?;

        // Traverse the cache using the path to (possibly) retrieve a value
        if let Some(ref mut val) = expr.get_mut(&mut self.cache) {
            **val = value;
        } else {
            expr.set(&mut self.cache, value);
        }
        Ok(())
    }

    pub fn get<'de, T: Deserialize<'de>>(&self, key: &str) -> Result<T> {
        // Parse the key into a path expression
        let expr: path::Expression = key.parse()?;

        // Traverse the cache using the path to (possibly) retrieve a value
        let value = expr.get(&self.cache).cloned();

        match value {
            Some(value) => {
                // Deserialize the received value into the requested type
                T::deserialize(value).map_err(|e| e.extend_with_key(key))
            }

            None => Err(ConfigError::NotFound(key.into())),
        }
    }

    pub fn get_str(&self, key: &str) -> Result<String> {
        self.get(key).and_then(Value::into_str)
    }

    pub fn get_int(&self, key: &str) -> Result<i64> {
        self.get(key).and_then(Value::into_int)
    }

    pub fn get_float(&self, key: &str) -> Result<f64> {
        self.get(key).and_then(Value::into_float)
    }

    pub fn get_bool(&self, key: &str) -> Result<bool> {
        self.get(key).and_then(Value::into_bool)
    }

    pub fn get_table(&self, key: &str) -> Result<HashMap<String, Value>> {
        self.get(key).and_then(Value::into_table)
    }

    pub fn get_array(&self, key: &str) -> Result<Vec<Value>> {
        self.get(key).and_then(Value::into_array)
    }

    /// Attempt to deserialize the entire configuration into the requested type.
    pub fn try_into<'de, T: Deserialize<'de>>(self) -> Result<T> {
        T::deserialize(self)
    }

    /// Attempt to serialize the entire configuration from the given type.
    pub fn try_from<T: Serialize>(from: &T) -> Result<Self> {
        let mut serializer = ConfigSerializer::default();
        from.serialize(&mut serializer)?;
        Ok(serializer.output)
    }

    #[deprecated(since = "0.7.0", note = "please use 'try_into' instead")]
    pub fn deserialize<'de, T: Deserialize<'de>>(self) -> Result<T> {
        self.try_into()
    }
}

impl Source for Config {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        self.cache.clone().into_table()
    }
}
<|endoftext|>