<commit_before>use std::rc::Rc;

use super::{config, Conditions, Message, TimerEvent};

use self::linear::LinearContext;
use self::map::MapContext;

#[derive(Debug)]
pub enum Context {
    Linear(LinearContext),
    Map(MapContext)
}

impl Context {
    pub fn on_timer(&mut self, event: &TimerEvent) {
        match *self {
            Context::Linear(ref mut context) => context.on_timer(event),
            Context::Map(ref mut context) => context.on_timer(event),
        }
    }

    pub fn on_message(&mut self, event: Rc<Message>) {
        match *self {
            Context::Linear(ref mut context) => context.on_message(event),
            Context::Map(ref mut context) => context.on_message(event),
        }
    }

    pub fn is_open(&mut self) -> bool {
        match *self {
            Context::Linear(ref context) => context.is_open(),
            Context::Map(ref mut context) => context.is_open(),
        }
    }

    pub fn new_linear(conditions: Conditions) -> Context {
        Context::Linear(
            LinearContext::new(conditions)
        )
    }

    pub fn new_map(conditions: Conditions) -> Context {
        Context::Map(
            MapContext::new(conditions)
        )
    }
}

impl From<config::Context> for Context {
    fn from(context: config::Context) -> Context {
        Context::new_linear(context.conditions)
    }
}

mod linear {
    use std::rc::Rc;

    use Conditions;
    use Message;
    use state::State;
    use TimerEvent;

    #[derive(Debug)]
    pub struct LinearContext {
        conditions: Conditions,
        state: State
    }

    impl LinearContext {
        pub fn new(conditions: Conditions) -> LinearContext {
            LinearContext {
                conditions: conditions,
                state: State::new()
            }
        }

        pub fn on_timer(&mut self, event: &TimerEvent) {
            self.conditions.on_timer(event, &mut self.state);
        }

        pub fn on_message(&mut self, event: Rc<Message>) {
            self.conditions.on_message(event, &mut self.state);
        }

        pub fn is_open(&self) -> bool {
            self.state.is_open()
        }
    }
}

mod map {
    use std::collections::BTreeMap;
    use std::fmt::Write;
    use std::rc::Rc;

    use Conditions;
    use Message;
    use state::State;
    use TimerEvent;

    #[derive(Debug)]
    pub struct MapContext {
        map: BTreeMap<String, State>,
        conditions: Conditions,
        format_buffer: String
    }

    impl MapContext {
        pub fn new(conditions: Conditions) -> MapContext {
            MapContext {
                map: BTreeMap::new(),
                conditions: conditions,
                format_buffer: String::new()
            }
        }

        pub fn on_timer(&mut self, event: &TimerEvent) {
            for (_, mut state) in self.map.iter_mut() {
                self.conditions.on_timer(event, &mut state);
            }
            self.remove_closed_states();
        }

        fn remove_closed_states(&mut self) {
            let ids_to_remove = self.map.iter().filter_map(|(id, state)| {
                if !state.is_open() {
                    Some(id.clone())
                } else {
                    None
                }
            }).collect::<Vec<String>>();

            for id in ids_to_remove {
                self.map.remove(&id);
            }
        }

        pub fn on_message(&mut self, event: Rc<Message>) {
            self.format_context_id(&event);
            let state = self.map.entry(self.format_buffer.clone()).or_insert(State::new());
            self.conditions.on_message(event, state);
            self.format_buffer.clear();
        }

        pub fn is_open(&mut self) -> bool {
            !self.map.is_empty()
        }

        fn format_context_id(&mut self, message: &Message) {
            let empty = String::new();
            let _ = self.format_buffer.write_str(message.get("HOST").unwrap_or(&empty));
            let _ = self.format_buffer.write_str(":");
            let _ = self.format_buffer.write_str(message.get("PROGRAM").unwrap_or(&empty));
            let _ = self.format_buffer.write_str(":");
            let _ = self.format_buffer.write_str(message.get("PID").unwrap_or(&empty));
        }
    }

    #[cfg(test)]
    mod test {
        use conditions::Builder;
        use Context;
        use TimerEvent;

        use std::rc::Rc;

        #[test]
        fn test_given_map_context_when_messages_have_the_same_kvpairs_then_they_go_to_the_same_context() {
            let delta = 10;
            let timeout = 30;
            let event = TimerEvent(delta);
            let patterns: Vec<String> = vec!["1".to_string(), "2".to_string(), "3".to_string()];
            let mut context = Context::new_map(Builder::new(timeout).patterns(patterns).build());
            let msg1 = Rc::new(btreemap! {
                "uuid".to_string() => "1".to_string(),
                "HOST".to_string() => "host".to_string(),
                "PROGRAM".to_string() => "program".to_string(),
                "PID".to_string() => "pid".to_string(),
            });
            let msg2 = Rc::new(btreemap! {
                "uuid".to_string() => "2".to_string(),
                "HOST".to_string() => "host2".to_string(),
                "PROGRAM".to_string() => "program2".to_string(),
                "PID".to_string() => "pid2".to_string(),
            });
            let msg3 = Rc::new(btreemap! {
                "uuid".to_string() => "3".to_string(),
                "HOST".to_string() => "host".to_string(),
                "PROGRAM".to_string() => "program".to_string(),
                "PID".to_string() => "pid".to_string(),
            });

            assert_false!(context.is_open());
            context.on_message(msg1.clone());
            assert_true!(context.is_open());
            context.on_timer(&event);
            context.on_message(msg2.clone());
            context.on_message(msg3.clone());
            context.on_timer(&event);
            context.on_timer(&event);
            assert_true!(context.is_open());
            context.on_timer(&event);
            assert_false!(context.is_open());
        }
    }
}

#[cfg(test)]
mod test {
    use std::rc::Rc;

    use TimerEvent;
    use super::Context;
    use conditions::Builder;

    #[test]
    fn test_given_close_condition_with_timeout_when_the_timeout_expires_then_the_condition_is_met() {
        let timeout = 100;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        println!("{:?}", &context);
        assert_false!(context.is_open());
        context.on_message(event);
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(50));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(49));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_false!(context.is_open());
    }

    #[test]
    fn test_given_close_condition_with_max_size_when_the_max_size_reached_then_the_condition_is_met() {
        let timeout = 100;
        let max_size = 3;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).max_size(max_size).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        println!("{:?}", &context);
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_message(event.clone());
        println!("{:?}", &context);
        assert_false!(context.is_open());
    }

    #[test]
    fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_without_renewing_messages_then_the_condition_is_met() {
        let timeout = 100;
        let renew_timeout = 10;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(8));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_false!(context.is_open());
    }

    #[test]
    fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_with_renewing_messages_then_the_context_is_not_closed() {
        let timeout = 100;
        let renew_timeout = 10;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        assert_false!(context.is_open());
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(8));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_true!(context.is_open());
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_true!(context.is_open());
    }
}
<commit_msg>close context on patterns<commit_after>use std::rc::Rc;

use super::{config, Conditions, Message, TimerEvent};

use self::linear::LinearContext;
use self::map::MapContext;

#[derive(Debug)]
pub enum Context {
    Linear(LinearContext),
    Map(MapContext)
}

impl Context {
    pub fn on_timer(&mut self, event: &TimerEvent) {
        match *self {
            Context::Linear(ref mut context) => context.on_timer(event),
            Context::Map(ref mut context) => context.on_timer(event),
        }
    }

    pub fn on_message(&mut self, event: Rc<Message>) {
        match *self {
            Context::Linear(ref mut context) => context.on_message(event),
            Context::Map(ref mut context) => context.on_message(event),
        }
    }

    pub fn is_open(&mut self) -> bool {
        match *self {
            Context::Linear(ref context) => context.is_open(),
            Context::Map(ref mut context) => context.is_open(),
        }
    }

    pub fn new_linear(conditions: Conditions) -> Context {
        Context::Linear(
            LinearContext::new(conditions)
        )
    }

    pub fn new_map(conditions: Conditions) -> Context {
        Context::Map(
            MapContext::new(conditions)
        )
    }
}

impl From<config::Context> for Context {
    fn from(context: config::Context) -> Context {
        Context::new_linear(context.conditions)
    }
}

mod linear {
    use std::rc::Rc;

    use Conditions;
    use Message;
    use state::State;
    use TimerEvent;

    #[derive(Debug)]
    pub struct LinearContext {
        conditions: Conditions,
        state: State
    }

    impl LinearContext {
        pub fn new(conditions: Conditions) -> LinearContext {
            LinearContext {
                conditions: conditions,
                state: State::new()
            }
        }

        pub fn on_timer(&mut self, event: &TimerEvent) {
            self.conditions.on_timer(event, &mut self.state);
        }

        pub fn on_message(&mut self, event: Rc<Message>) {
            self.conditions.on_message(event, &mut self.state);
        }

        pub fn is_open(&self) -> bool {
            self.state.is_open()
        }
    }
}

mod map {
    use std::collections::BTreeMap;
    use std::fmt::Write;
    use std::rc::Rc;

    use Conditions;
    use Message;
    use state::State;
    use TimerEvent;

    #[derive(Debug)]
    pub struct MapContext {
        map: BTreeMap<String, State>,
        conditions: Conditions,
        format_buffer: String
    }

    impl MapContext {
        pub fn new(conditions: Conditions) -> MapContext {
            MapContext {
                map: BTreeMap::new(),
                conditions: conditions,
                format_buffer: String::new()
            }
        }

        pub fn on_timer(&mut self, event: &TimerEvent) {
            for (_, mut state) in self.map.iter_mut() {
                self.conditions.on_timer(event, &mut state);
            }
            self.remove_closed_states();
        }

        fn remove_closed_states(&mut self) {
            let ids_to_remove = self.map.iter().filter_map(|(id, state)| {
                if !state.is_open() {
                    Some(id.clone())
                } else {
                    None
                }
            }).collect::<Vec<String>>();

            for id in ids_to_remove {
                self.map.remove(&id);
            }
        }

        pub fn on_message(&mut self, event: Rc<Message>) {
            self.format_context_id(&event);
            self.update_state(event);
            self.format_buffer.clear();
        }

        fn update_state(&mut self, event: Rc<Message>) {
            let id = self.format_buffer.clone();

            match self.map.remove(&id) {
                Some(mut state) => {
                    self.conditions.on_message(event, &mut state);
                    if state.is_open() {
                        self.map.insert(id, state);
                    }
                },
                None => {
                    let mut state = State::new();
                    self.conditions.on_message(event, &mut state);
                    self.map.insert(id, state);
                }
            }
        }

        pub fn is_open(&mut self) -> bool {
            !self.map.is_empty()
        }

        fn format_context_id(&mut self, message: &Message) {
            let empty = String::new();
            let _ = self.format_buffer.write_str(message.get("HOST").unwrap_or(&empty));
            let _ = self.format_buffer.write_str(":");
            let _ = self.format_buffer.write_str(message.get("PROGRAM").unwrap_or(&empty));
            let _ = self.format_buffer.write_str(":");
            let _ = self.format_buffer.write_str(message.get("PID").unwrap_or(&empty));
        }
    }

    #[cfg(test)]
    mod test {
        use conditions::Builder;
        use Context;
        use TimerEvent;

        use std::rc::Rc;

        #[test]
        fn test_given_map_context_when_messages_have_the_same_kvpairs_then_they_go_to_the_same_context() {
            let delta = 10;
            let timeout = 30;
            let event = TimerEvent(delta);
            let patterns: Vec<String> = vec!["1".to_string(), "2".to_string(), "3".to_string()];
            let mut context = Context::new_map(Builder::new(timeout).patterns(patterns).build());
            let msg1 = Rc::new(btreemap! {
                "uuid".to_string() => "1".to_string(),
                "HOST".to_string() => "host".to_string(),
                "PROGRAM".to_string() => "program".to_string(),
                "PID".to_string() => "pid".to_string(),
            });
            let msg2 = Rc::new(btreemap! {
                "uuid".to_string() => "2".to_string(),
                "HOST".to_string() => "host2".to_string(),
                "PROGRAM".to_string() => "program2".to_string(),
                "PID".to_string() => "pid2".to_string(),
            });
            let msg3 = Rc::new(btreemap! {
                "uuid".to_string() => "3".to_string(),
                "HOST".to_string() => "host".to_string(),
                "PROGRAM".to_string() => "program".to_string(),
                "PID".to_string() => "pid".to_string(),
            });

            assert_false!(context.is_open());
            context.on_message(msg1.clone());
            assert_true!(context.is_open());
            context.on_timer(&event);
            context.on_message(msg2.clone());
            context.on_message(msg3.clone());
            context.on_timer(&event);
            context.on_timer(&event);
            assert_true!(context.is_open());
            context.on_timer(&event);
            assert_false!(context.is_open());
        }
    }
}

#[cfg(test)]
mod test {
    use std::rc::Rc;

    use TimerEvent;
    use super::Context;
    use conditions::Builder;

    #[test]
    fn test_given_close_condition_with_timeout_when_the_timeout_expires_then_the_condition_is_met() {
        let timeout = 100;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        println!("{:?}", &context);
        assert_false!(context.is_open());
        context.on_message(event);
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(50));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(49));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_false!(context.is_open());
    }

    #[test]
    fn test_given_close_condition_with_max_size_when_the_max_size_reached_then_the_condition_is_met() {
        let timeout = 100;
        let max_size = 3;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).max_size(max_size).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        println!("{:?}", &context);
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_message(event.clone());
        println!("{:?}", &context);
        assert_false!(context.is_open());
    }

    #[test]
    fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_without_renewing_messages_then_the_condition_is_met() {
        let timeout = 100;
        let renew_timeout = 10;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(8));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_false!(context.is_open());
    }

    #[test]
    fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_with_renewing_messages_then_the_context_is_not_closed() {
        let timeout = 100;
        let renew_timeout = 10;
        let msg_id = "1".to_string();
        let mut context = Context::new_linear(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
        let msg1 = btreemap!{
            "uuid".to_string() => msg_id.clone(),
        };
        let event = Rc::new(msg1);
        assert_false!(context.is_open());
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(8));
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_true!(context.is_open());
        context.on_message(event.clone());
        assert_true!(context.is_open());
        context.on_timer(&mut TimerEvent(1));
        assert_true!(context.is_open());
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Stdout, stdout, stdin};
use termion::input::TermRead;
use termion::raw::{IntoRawMode, RawTerminal};

use super::*;

/// The default for `Context.word_fn`.
pub fn get_buffer_words(buf: &Buffer) -> Vec<(usize, usize)> {
    let mut res = Vec::new();

    let mut word_start = None;
    let mut just_had_backslash = false;

    for (i, &c) in buf.chars().enumerate() {
        if c == '\\' {
            just_had_backslash = true;
            continue;
        }

        if let Some(start) = word_start {
            if c == ' ' && !just_had_backslash {
                res.push((start, i));
                word_start = None;
            }
        } else {
            if c != ' ' {
                word_start = Some(i);
            }
        }

        just_had_backslash = false;
    }

    if let Some(start) = word_start {
        res.push((start, buf.num_chars()));
    }

    res
}

pub struct Context {
    pub history: History,
    pub completer: Option<Box<Completer>>,
    pub word_fn: Box<Fn(&Buffer) -> Vec<(usize, usize)>>,
}

impl Context {
    pub fn new() -> Self {
        Context {
            history: History::new(),
            completer: None,
            word_fn: Box::new(get_buffer_words),
        }
    }

    /// Creates an `Editor` and feeds it keypresses from stdin until the line is entered.
    /// The output is stdout.
    /// The returned line has the newline removed.
    /// Before returning, will revert all changes to the history buffers.
    pub fn read_line<P: Into<String>>(&mut self,
                                      prompt: P,
                                      mut handler: &mut EventHandler<RawTerminal<Stdout>>)
                                      -> io::Result<String> {
        let res = {
            let stdout = stdout().into_raw_mode().unwrap();
            let ed = try!(Editor::new(stdout, prompt.into(), self));
            Self::handle_keys(ed, handler)
        };

        self.revert_all_history();
        res
    }

    fn handle_keys(mut ed: Editor<RawTerminal<Stdout>>,
                   mut handler: &mut EventHandler<RawTerminal<Stdout>>)
    -> io::Result<String>
    {
        let stdin = stdin();
        for c in stdin.keys() {
            if try!(ed.handle_key(c.unwrap(), handler)) {
                break;
            }
        }

        Ok(ed.into())
    }

    pub fn revert_all_history(&mut self) {
        for buf in &mut self.history.buffers {
            buf.revert();
        }
    }
}
<commit_msg>use KeyMap instead of Editor in Context<commit_after>use std::io::{self, Stdout, stdout, stdin, Write};
use termion::input::TermRead;
use termion::raw::{IntoRawMode, RawTerminal};

use super::*;

/// The default for `Context.word_fn`.
pub fn get_buffer_words(buf: &Buffer) -> Vec<(usize, usize)> {
    let mut res = Vec::new();

    let mut word_start = None;
    let mut just_had_backslash = false;

    for (i, &c) in buf.chars().enumerate() {
        if c == '\\' {
            just_had_backslash = true;
            continue;
        }

        if let Some(start) = word_start {
            if c == ' ' && !just_had_backslash {
                res.push((start, i));
                word_start = None;
            }
        } else {
            if c != ' ' {
                word_start = Some(i);
            }
        }

        just_had_backslash = false;
    }

    if let Some(start) = word_start {
        res.push((start, buf.num_chars()));
    }

    res
}

pub struct Context {
    pub history: History,
    pub completer: Option<Box<Completer>>,
    pub word_fn: Box<Fn(&Buffer) -> Vec<(usize, usize)>>,
}

impl Context {
    pub fn new() -> Self {
        Context {
            history: History::new(),
            completer: None,
            word_fn: Box::new(get_buffer_words),
        }
    }

    /// Creates an `Editor` and feeds it keypresses from stdin until the line is entered.
    /// The output is stdout.
    /// The returned line has the newline removed.
    /// Before returning, will revert all changes to the history buffers.
    pub fn read_line<P: Into<String>>(&mut self,
                                      prompt: P,
                                      mut handler: &mut EventHandler<RawTerminal<Stdout>>)
                                      -> io::Result<String> {
        let res = {
            let stdout = stdout().into_raw_mode().unwrap();
            let ed = try!(Editor::new(stdout, prompt.into(), self));
            Self::handle_keys(ed, handler)
        };

        self.revert_all_history();
        res
    }

    fn handle_keys<'a, T, W: Write, M: KeyMap<'a, W, T>>(
        mut keymap: M,
        mut handler: &mut EventHandler<W>)
    -> io::Result<String>
        where String: From<M>
    {
        let stdin = stdin();
        for c in stdin.keys() {
            if try!(keymap.handle_key(c.unwrap(), handler)) {
                break;
            }
        }

        Ok(keymap.into())
    }

    pub fn revert_all_history(&mut self) {
        for buf in &mut self.history.buffers {
            buf.revert();
        }
    }
}
<|endoftext|>


<commit_before>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::EnvironmentProvider;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    let rows = conn.query("SELECT path, mime, date_updated, content
                           FROM files
                           WHERE path = $1", &[&path]).unwrap();

    if rows.len() == 0 {
        None
    } else {
        let row = rows.get(0);
        let mut content = row.get(3);
        if content == b"in-s3" {
            let client = s3_client();
            content = client.and_then(|c| c.get_object(GetObjectRequest {
                bucket: "rust-docs-rs".into(),
                key: path.into(),
                ..Default::default()
            }).sync().ok()).and_then(|r| r.body).map(|b| {
                let mut b = b.into_blocking_read();
                let mut content = Vec::new();
                b.read_to_end(&mut content).unwrap();
                content
            }).unwrap();
        };

        Some(Blob {
            path: row.get(0),
            mime: row.get(1),
            date_updated: row.get(2),
            content,
        })
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() {
        return None;
    }
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        EnvironmentProvider::default(),
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());
    let client = s3_client();
    let mut file_list_with_mimes: Vec<(String, PathBuf)> = Vec::new();

    for file_path in try!(get_file_list(&path)) {
        let (path, content, mime) = {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            let content: Option<Vec<u8>> = if let Some(client) = &client {
                let s3_res = client.put_object(PutObjectRequest {
                    bucket: "rust-docs-rs".into(),
                    key: bucket_path.clone(),
                    body: Some(content.clone().into()),
                    content_type: Some(mime.clone()),
                    ..Default::default()
                }).sync();
                match s3_res {
                    // we've successfully uploaded the content, so steal it;
                    // we don't want to put it in the DB
                    Ok(_) => None,
                    // Since s3 was configured, we want to panic on failure to upload.
                    Err(e) => {
                        panic!("failed to upload to {}: {:?}", bucket_path, e)
                    },
                }
            } else {
                Some(content.clone().into())
            };

            file_list_with_mimes.push((mime.clone(), file_path.clone()));

            (
                bucket_path,
                content,
                mime,
            )
        };

        // check if file already exists in database
        let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&path]));

        let content = content.unwrap_or_else(|| "in-s3".to_owned().into());

        if rows.get(0).get::<usize, i64>(0) == 0 {
            try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                             &[&path, &mime, &content]));
        } else {
            try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                              WHERE path = $1",
                             &[&path, &mime, &content]));
        }
    }

    try!(trans.commit());

    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::current_thread::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("UPDATE files SET content = E'in-s3' WHERE path = $1")
                .unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<commit_msg>Fix uploads of large crates to S3 breaking the client<commit_after>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::{PathBuf, Path};
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::fs;
use std::io::Read;
use error::Result;
use failure::err_msg;
use rusoto_s3::{S3, PutObjectRequest, GetObjectRequest, S3Client};
use rusoto_core::region::Region;
use rusoto_credential::EnvironmentProvider;

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          files: &mut Vec<PathBuf>)
                                          -> Result<()> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            files.push(file.path());
        } else if try!(file.file_type()).is_dir() {
            try!(get_file_list_from_dir(file.path(), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
    let path = path.as_ref();
    let mut files = Vec::new();

    if !path.exists() {
        return Err(err_msg("File not found"));
    } else if path.is_file() {
        files.push(PathBuf::from(path.file_name().unwrap()));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, &mut files));
        for file_path in &mut files {
            // We want the paths in this list to not be {path}/bar.txt but just bar.txt
            *file_path = PathBuf::from(file_path.strip_prefix(path).unwrap());
        }
    }

    Ok(files)
}

pub struct Blob {
    pub path: String,
    pub mime: String,
    pub date_updated: time::Timespec,
    pub content: Vec<u8>,
}

pub fn get_path(conn: &Connection, path: &str) -> Option<Blob> {
    let rows = conn.query("SELECT path, mime, date_updated, content
                           FROM files
                           WHERE path = $1", &[&path]).unwrap();

    if rows.len() == 0 {
        None
    } else {
        let row = rows.get(0);
        let mut content = row.get(3);
        if content == b"in-s3" {
            let client = s3_client();
            content = client.and_then(|c| c.get_object(GetObjectRequest {
                bucket: "rust-docs-rs".into(),
                key: path.into(),
                ..Default::default()
            }).sync().ok()).and_then(|r| r.body).map(|b| {
                let mut b = b.into_blocking_read();
                let mut content = Vec::new();
                b.read_to_end(&mut content).unwrap();
                content
            }).unwrap();
        };

        Some(Blob {
            path: row.get(0),
            mime: row.get(1),
            date_updated: row.get(2),
            content,
        })
    }
}

fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() {
        return None;
    }
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        EnvironmentProvider::default(),
        std::env::var("S3_ENDPOINT").ok().map(|e| Region::Custom {
            name: "us-west-1".to_owned(),
            endpoint: e,
        }).unwrap_or(Region::UsWest1),
    ))
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    try!(cookie.load::<&str>(&[]));

    let trans = try!(conn.transaction());
    let mut client = s3_client();
    let mut file_list_with_mimes: Vec<(String, PathBuf)> = Vec::new();

    for file_path in try!(get_file_list(&path)) {
        let (path, content, mime) = {
            let path = Path::new(path.as_ref()).join(&file_path);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match fs::File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let bucket_path = Path::new(prefix).join(&file_path)
                .into_os_string().into_string().unwrap();

            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                // do the same for javascript files
                if mime == "text/plain" {
                    let e = file_path.extension().unwrap_or_default();
                    if e == "css" {
                        "text/css".to_owned()
                    } else if e == "js" {
                        "application/javascript".to_owned()
                    } else {
                        mime.to_owned()
                    }
                } else {
                    mime.to_owned()
                }
            };

            let content: Option<Vec<u8>> = if let Some(client) = &mut client {
                let mut attempts = 0;
                loop {
                    let s3_res = client.put_object(PutObjectRequest {
                        bucket: "rust-docs-rs".into(),
                        key: bucket_path.clone(),
                        body: Some(content.clone().into()),
                        content_type: Some(mime.clone()),
                        ..Default::default()
                    }).sync();
                    attempts += 1;
                    match s3_res {
                        // we've successfully uploaded the content, so steal it;
                        // we don't want to put it in the DB
                        Ok(_) => break None,
                        // Since s3 was configured, we want to panic on failure to upload.
                        Err(e) => {
                            log::error!("failed to upload to {}: {:?}", bucket_path, e);
                            // Get a new client, in case the old one's connection is stale.
                            // AWS will kill our connection if it's alive for too long; this avoids
                            // that preventing us from building the crate entirely.
                            *client = s3_client().unwrap();
                            if attempts > 3 {
                                panic!("failed to upload 3 times, exiting");
                            } else {
                                continue;
                            }
                        },
                    }
                }
            } else {
                Some(content.clone().into())
            };

            file_list_with_mimes.push((mime.clone(), file_path.clone()));

            (
                bucket_path,
                content,
                mime,
            )
        };

        // check if file already exists in database
        let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&path]));

        let content = content.unwrap_or_else(|| "in-s3".to_owned().into());

        if rows.get(0).get::<usize, i64>(0) == 0 {
            try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                             &[&path, &mime, &content]));
        } else {
            try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                              WHERE path = $1",
                             &[&path, &mime, &content]));
        }
    }

    try!(trans.commit());

    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, PathBuf)>) -> Result<Json> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.into_os_string().into_string().unwrap());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize> {
    let trans = try!(conn.transaction());
    let client = s3_client().expect("configured s3");

    let rows = try!(trans.query(
            &format!("SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}", n),
            &[]));
    let count = rows.len();

    let mut rt = ::tokio::runtime::current_thread::Runtime::new().unwrap();
    let mut futures = Vec::new();
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(client.put_object(PutObjectRequest {
            bucket: "rust-docs-rs".into(),
            key: path.clone(),
            body: Some(content.into()),
            content_type: Some(mime),
            ..Default::default()
        }).map(move |_| {
            path_1
        }).map_err(move |e| {
            panic!("failed to upload to {}: {:?}", path, e)
        }));
    }

    use ::futures::future::Future;
    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("UPDATE files SET content = E'in-s3' WHERE path = $1")
                .unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    try!(trans.commit());

    Ok(count)
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::get_file_list;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::try_init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], std::path::Path::new("Cargo.toml"));
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The Rust Project Developers.
// Copyright 2018 The Uuid Project Developers.
//
// See the COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::error;

<<<<<<< HEAD
impl error::Error for crate::BytesError {
=======
impl error::Error for super::BytesError {
>>>>>>> master
    fn description(&self) -> &str {
        "invalid number of uuid bytes"
    }
}

impl error::Error for super::Error {
    fn description(&self) -> &str {
        match *self {
            super::Error::Bytes(ref err) => error::Error::description(err),
            super::Error::Parse(ref err) => error::Error::description(err),
        }
    }
}
<commit_msg>fix missed merge conflict<commit_after>// Copyright 2013-2014 The Rust Project Developers.
// Copyright 2018 The Uuid Project Developers.
//
// See the COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::error;

impl error::Error for crate::BytesError {
    fn description(&self) -> &str {
        "invalid number of uuid bytes"
    }
}

impl error::Error for super::Error {
    fn description(&self) -> &str {
        match *self {
            super::Error::Bytes(ref err) => error::Error::description(err),
            super::Error::Parse(ref err) => error::Error::description(err),
        }
    }
}
<|endoftext|>


<commit_before>//! Indirect system call
//!
use libc::c_int;

pub use self::arch::*;

#[cfg(target_arch = "x86_64")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 155;
    pub static MEMFD_CREATE: Syscall = 319;
}

#[cfg(target_arch = "x86")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 217;
    pub static MEMFD_CREATE: Syscall = 356;
}

#[cfg(target_arch = "aarch64")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 41;
    pub static MEMFD_CREATE: Syscall = 279;
}

#[cfg(target_arch = "arm")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 218;
    pub static MEMFD_CREATE: Syscall = 385;
}

#[cfg(target_arch = "mips")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 216;
    pub static MEMFD_CREATE: Syscall = 354;
}

extern {
    pub fn syscall(num: Syscall, ...) -> c_int;
}
<commit_msg>Added missing syscalls for powerpc<commit_after>//! Indirect system call
//!
use libc::c_int;

pub use self::arch::*;

#[cfg(target_arch = "x86_64")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 155;
    pub static MEMFD_CREATE: Syscall = 319;
}

#[cfg(target_arch = "x86")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 217;
    pub static MEMFD_CREATE: Syscall = 356;
}

#[cfg(target_arch = "aarch64")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 41;
    pub static MEMFD_CREATE: Syscall = 279;
}

#[cfg(target_arch = "arm")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 218;
    pub static MEMFD_CREATE: Syscall = 385;
}

#[cfg(target_arch = "mips")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 216;
    pub static MEMFD_CREATE: Syscall = 354;
}

#[cfg(target_arch = "powerpc")]
mod arch {
    use libc::c_long;

    pub type Syscall = c_long;

    pub static SYSPIVOTROOT: Syscall = 203;
    pub static MEMFD_CREATE: Syscall = 360;
}

extern {
    pub fn syscall(num: Syscall, ...) -> c_int;
}
<|endoftext|>


<commit_before>pub mod info;
pub mod peer;
mod picker;

pub mod piece_field;

pub use self::piece_field::PieceField;
pub use self::info::Info;
pub use self::peer::Peer;

use bencode::BEncode;
use self::peer::Message;
use self::picker::Picker;
use std::{fmt, io, cmp};
use {disk, DISK, tracker, TRACKER};
use pbr::ProgressBar;
use std::collections::HashSet;
use util::io_err;

pub struct Torrent {
    pub info: Info,
    pub pieces: PieceField,
    pub uploaded: usize,
    pub downloaded: usize,
    pub id: usize,
    peers: HashSet<usize>,
    picker: Picker,
    pb: ProgressBar<io::Stdout>,
}

impl fmt::Debug for Torrent {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Torrent {{ info: {:?} }}", self.info)
    }
}

pub struct Broadcast {
    pub msg: Message,
    pub peers: HashSet<usize>,
}

impl Torrent {
    pub fn from_bencode(data: BEncode) -> Result<Torrent, &'static str> {
        let info = Info::from_bencode(data)?;
        println!("Handling with torrent with {:?} pl, {:?} pieces, {:?} sf len", info.piece_len, info.pieces(), info.files.last().unwrap().length);
        // Create dummy files
        info.create_files().unwrap();
        let peers = HashSet::new();
        let pieces = PieceField::new(info.pieces());
        let picker = Picker::new(&info);
        let pb = ProgressBar::new(info.pieces() as u64);
        Ok(Torrent { id: 0, info, peers, pieces, picker, pb, uploaded: 0, downloaded: 0 })
    }

    pub fn block_available(&mut self, peer: &mut Peer, resp: disk::Response) -> io::Result<()> {
        let ctx = resp.context;
        let p = Message::s_piece(ctx.idx, ctx.begin, ctx.length, resp.data);
        peer.send_message(p)?;
        Ok(())
    }

    pub fn peer_readable(&mut self, peer: &mut Peer) -> io::Result<Vec<Broadcast>> {
        let mut resp = Vec::new();
        let res = peer.readable()?;
        for msg in res {
            resp.extend(self.handle_msg(msg, peer)?);
        }
        Ok(resp)
    }

    pub fn handle_msg(&mut self, msg: Message, peer: &mut Peer) -> io::Result<Vec<Broadcast>> {
        let mut v = Vec::with_capacity(1);
        match msg {
            Message::Handshake { .. } => {
                println!("Connection established with peer {:?}", peer.id);
            }
            Message::Bitfield(mut pf) => {
                pf.cap(self.pieces.len());
                peer.pieces = pf;
                if self.pieces.usable(&peer.pieces) {
                    peer.send_message(Message::Interested)?;
                }
            }
            Message::Have(idx) => {
                peer.pieces.set_piece(idx);
            }
            Message::Unchoke => {
                peer.being_choked = false;
                Torrent::make_requests(&mut self.picker, peer, &self.info)?;
            }
            Message::Choke => {
                peer.being_choked = true;
            }
            Message::Piece { index, begin, data, length } => {
                if self.pieces.complete() {
                    return Ok(v);
                }

                peer.queued -= 1;
                Torrent::write_piece(&self.info, index, begin, length, data);
                let (piece_done, mut peers) = self.picker.completed(index, begin);
                if piece_done {
                    self.downloaded += 1;
                    self.pb.inc();
                    self.pieces.set_piece(index);
                    if self.pieces.complete() {
                        TRACKER.tx.send(tracker::Request::completed(&self)).unwrap();
                        self.pb.finish_print("Downloaded!");
                    }
                    let m = Message::Have(index);
                    v.push(Broadcast { msg: m, peers: self.peers.clone() });
                }
                if peers.len() > 1 {
                    peers.remove(&peer.id);
                    let m = Message::Cancel { index, begin, length };
                    v.push(Broadcast { msg: m, peers: self.peers.clone() });
                }
                if !peer.being_choked && !self.pieces.complete() {
                    Torrent::make_requests(&mut self.picker, peer, &self.info)?;
                }
            }
            Message::Request { index, begin, length } => {
                // TODO get this from some sort of allocator.
                if !peer.choked {
                    Torrent::request_read(peer.id, &self.info, index, begin, length, Box::new([0u8; 16384]));
                } else {
                    return io_err("Peer requested while choked!");
                }
            }
            Message::Interested => {
                peer.interested = true;
                // If we're in seed mode, upload, otherwise
                // use the general tit for tat heuristic for unchoking.
                // TODO: implement said heuristic
                if self.pieces.complete() {
                    peer.choked = false;
                    peer.send_message(Message::Unchoke)?;
                }
            }
            Message::Uninterested => {
                peer.interested = false;
                if !peer.choked {
                    peer.choked = true;
                    peer.send_message(Message::Choke)?;
                }
            }
            _ => { }
        }
        Ok(v)
    }

    /// Calculates the file offsets for a given index, begin, and block length.
    fn calc_block_locs(info: &Info, index: u32, begin: u32, mut len: u32) -> Vec<disk::Location> {
        // The absolute byte offset where we start processing data.
        let mut cur_start = index * info.piece_len as u32 + begin;
        // Current index of the data block we're writing
        let mut data_start = 0;
        // The current file end length.
        let mut fidx = 0;
        // Iterate over all file lengths, if we find any which end a bigger
        // idx than cur_start, write from cur_start..cur_start + file_write_len for that file
        // and continue if we're now at the end of the file.
        let mut locs = Vec::new();
        for f in info.files.iter() {
            fidx += f.length;
            if (cur_start as usize) < fidx {
                let file_write_len = cmp::min(fidx - cur_start as usize, len as usize);
                let offset = (cur_start - (fidx - f.length) as u32) as u64;
                if file_write_len == len as usize {
                    // The file is longer than our len, just write to it,
                    // exit loop
                    locs.push(disk::Location::new(f.path.clone(), offset, data_start, data_start + file_write_len));
                    break;
                } else {
                    // Write to the end of file, continue
                    locs.push(disk::Location::new(f.path.clone(), offset, data_start, data_start + file_write_len as usize));
                    len -= file_write_len as u32;
                    cur_start += file_write_len as u32;
                    data_start += file_write_len;
                }
            }
        }
        locs
    }

    #[inline(always)]
    /// Writes a piece of torrent info, with piece index idx,
    /// piece offset begin, piece length of len, and data bytes.
    /// The disk send handle is also provided.
    fn write_piece(info: &Info, index: u32, begin: u32, len: u32, data: Box<[u8; 16384]>) {
        let locs = Torrent::calc_block_locs(info, index, begin, len);
        DISK.tx.send(disk::Request::write(data, locs)).unwrap();
    }

    #[inline(always)]
    /// Issues a read request of the given torrent
    fn request_read(id: usize, info: &Info, index: u32, begin: u32, len: u32, data: Box<[u8; 16384]>) {
        let locs = Torrent::calc_block_locs(info, index, begin, len);
        let ctx = disk::Ctx::new(id, index, begin, len);
        DISK.tx.send(disk::Request::read(ctx, data, locs)).unwrap();
    }

    #[inline(always)]
    fn make_requests(picker: &mut Picker, peer: &mut Peer, info: &Info) -> io::Result<()> {
        // keep 5 outstanding requests?
        while peer.queued < 5 {
            if let Some((idx, offset)) = picker.pick(&peer) {
                if info.is_last_piece((idx, offset)) {
                    peer.send_message(Message::request(idx, offset, info.last_piece_len()))?;
                } else {
                    peer.send_message(Message::request(idx, offset, 16384))?;
                }
                peer.queued += 1;
            } else {
                break;
            }
        }
        Ok(())
    }

    pub fn peer_writable(&mut self, peer: &mut Peer) -> io::Result<bool> {
        peer.writable()
    }

    pub fn file_size(&self) -> usize {
        let mut size = 0;
        for file in self.info.files.iter() {
            size += file.length;
        }
        size
    }

    pub fn add_peer(&mut self, peer: &mut Peer) -> io::Result<()> {
        peer.set_torrent(&self)?;
        self.peers.insert(peer.id);
        Ok(())
    }

    pub fn remove_peer(&mut self, peer: &mut Peer) {
        self.peers.remove(&peer.id);
    }
}
<commit_msg>Proper broadcast cancellation<commit_after>pub mod info;
pub mod peer;
mod picker;

pub mod piece_field;

pub use self::piece_field::PieceField;
pub use self::info::Info;
pub use self::peer::Peer;

use bencode::BEncode;
use self::peer::Message;
use self::picker::Picker;
use std::{fmt, io, cmp};
use {disk, DISK, tracker, TRACKER};
use pbr::ProgressBar;
use std::collections::HashSet;
use util::io_err;

pub struct Torrent {
    pub info: Info,
    pub pieces: PieceField,
    pub uploaded: usize,
    pub downloaded: usize,
    pub id: usize,
    peers: HashSet<usize>,
    picker: Picker,
    pb: ProgressBar<io::Stdout>,
}

impl fmt::Debug for Torrent {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Torrent {{ info: {:?} }}", self.info)
    }
}

pub struct Broadcast {
    pub msg: Message,
    pub peers: HashSet<usize>,
}

impl Torrent {
    pub fn from_bencode(data: BEncode) -> Result<Torrent, &'static str> {
        let info = Info::from_bencode(data)?;
        println!("Handling with torrent with {:?} pl, {:?} pieces, {:?} sf len", info.piece_len, info.pieces(), info.files.last().unwrap().length);
        // Create dummy files
        info.create_files().unwrap();
        let peers = HashSet::new();
        let pieces = PieceField::new(info.pieces());
        let picker = Picker::new(&info);
        let pb = ProgressBar::new(info.pieces() as u64);
        Ok(Torrent { id: 0, info, peers, pieces, picker, pb, uploaded: 0, downloaded: 0 })
    }

    pub fn block_available(&mut self, peer: &mut Peer, resp: disk::Response) -> io::Result<()> {
        let ctx = resp.context;
        let p = Message::s_piece(ctx.idx, ctx.begin, ctx.length, resp.data);
        peer.send_message(p)?;
        Ok(())
    }

    pub fn peer_readable(&mut self, peer: &mut Peer) -> io::Result<Vec<Broadcast>> {
        let mut resp = Vec::new();
        let res = peer.readable()?;
        for msg in res {
            resp.extend(self.handle_msg(msg, peer)?);
        }
        Ok(resp)
    }

    pub fn handle_msg(&mut self, msg: Message, peer: &mut Peer) -> io::Result<Vec<Broadcast>> {
        let mut v = Vec::with_capacity(1);
        match msg {
            Message::Handshake { .. } => {
                println!("Connection established with peer {:?}", peer.id);
            }
            Message::Bitfield(mut pf) => {
                pf.cap(self.pieces.len());
                peer.pieces = pf;
                if self.pieces.usable(&peer.pieces) {
                    peer.send_message(Message::Interested)?;
                }
            }
            Message::Have(idx) => {
                peer.pieces.set_piece(idx);
            }
            Message::Unchoke => {
                peer.being_choked = false;
                Torrent::make_requests(&mut self.picker, peer, &self.info)?;
            }
            Message::Choke => {
                peer.being_choked = true;
            }
            Message::Piece { index, begin, data, length } => {
                if self.pieces.complete() {
                    return Ok(v);
                }

                peer.queued -= 1;
                Torrent::write_piece(&self.info, index, begin, length, data);
                let (piece_done, mut peers) = self.picker.completed(index, begin);
                if piece_done {
                    self.downloaded += 1;
                    self.pb.inc();
                    self.pieces.set_piece(index);
                    if self.pieces.complete() {
                        TRACKER.tx.send(tracker::Request::completed(&self)).unwrap();
                        self.pb.finish_print("Downloaded!");
                    }
                    let m = Message::Have(index);
                    v.push(Broadcast { msg: m, peers: self.peers.clone() });
                }
                if peers.len() > 1 {
                    peers.remove(&peer.id);
                    let m = Message::Cancel { index, begin, length };
                    v.push(Broadcast { msg: m, peers });
                }
                if !peer.being_choked && !self.pieces.complete() {
                    Torrent::make_requests(&mut self.picker, peer, &self.info)?;
                }
            }
            Message::Request { index, begin, length } => {
                // TODO get this from some sort of allocator.
                if !peer.choked {
                    Torrent::request_read(peer.id, &self.info, index, begin, length, Box::new([0u8; 16384]));
                } else {
                    return io_err("Peer requested while choked!");
                }
            }
            Message::Interested => {
                peer.interested = true;
                // If we're in seed mode, upload, otherwise
                // use the general tit for tat heuristic for unchoking.
                // TODO: implement said heuristic
                if self.pieces.complete() {
                    peer.choked = false;
                    peer.send_message(Message::Unchoke)?;
                }
            }
            Message::Uninterested => {
                peer.interested = false;
                if !peer.choked {
                    peer.choked = true;
                    peer.send_message(Message::Choke)?;
                }
            }
            _ => { }
        }
        Ok(v)
    }

    /// Calculates the file offsets for a given index, begin, and block length.
    fn calc_block_locs(info: &Info, index: u32, begin: u32, mut len: u32) -> Vec<disk::Location> {
        // The absolute byte offset where we start processing data.
        let mut cur_start = index * info.piece_len as u32 + begin;
        // Current index of the data block we're writing
        let mut data_start = 0;
        // The current file end length.
        let mut fidx = 0;
        // Iterate over all file lengths, if we find any which end a bigger
        // idx than cur_start, write from cur_start..cur_start + file_write_len for that file
        // and continue if we're now at the end of the file.
        let mut locs = Vec::new();
        for f in info.files.iter() {
            fidx += f.length;
            if (cur_start as usize) < fidx {
                let file_write_len = cmp::min(fidx - cur_start as usize, len as usize);
                let offset = (cur_start - (fidx - f.length) as u32) as u64;
                if file_write_len == len as usize {
                    // The file is longer than our len, just write to it,
                    // exit loop
                    locs.push(disk::Location::new(f.path.clone(), offset, data_start, data_start + file_write_len));
                    break;
                } else {
                    // Write to the end of file, continue
                    locs.push(disk::Location::new(f.path.clone(), offset, data_start, data_start + file_write_len as usize));
                    len -= file_write_len as u32;
                    cur_start += file_write_len as u32;
                    data_start += file_write_len;
                }
            }
        }
        locs
    }

    #[inline(always)]
    /// Writes a piece of torrent info, with piece index idx,
    /// piece offset begin, piece length of len, and data bytes.
    /// The disk send handle is also provided.
    fn write_piece(info: &Info, index: u32, begin: u32, len: u32, data: Box<[u8; 16384]>) {
        let locs = Torrent::calc_block_locs(info, index, begin, len);
        DISK.tx.send(disk::Request::write(data, locs)).unwrap();
    }

    #[inline(always)]
    /// Issues a read request of the given torrent
    fn request_read(id: usize, info: &Info, index: u32, begin: u32, len: u32, data: Box<[u8; 16384]>) {
        let locs = Torrent::calc_block_locs(info, index, begin, len);
        let ctx = disk::Ctx::new(id, index, begin, len);
        DISK.tx.send(disk::Request::read(ctx, data, locs)).unwrap();
    }

    #[inline(always)]
    fn make_requests(picker: &mut Picker, peer: &mut Peer, info: &Info) -> io::Result<()> {
        // keep 5 outstanding requests?
        while peer.queued < 5 {
            if let Some((idx, offset)) = picker.pick(&peer) {
                if info.is_last_piece((idx, offset)) {
                    peer.send_message(Message::request(idx, offset, info.last_piece_len()))?;
                } else {
                    peer.send_message(Message::request(idx, offset, 16384))?;
                }
                peer.queued += 1;
            } else {
                break;
            }
        }
        Ok(())
    }

    pub fn peer_writable(&mut self, peer: &mut Peer) -> io::Result<bool> {
        peer.writable()
    }

    pub fn file_size(&self) -> usize {
        let mut size = 0;
        for file in self.info.files.iter() {
            size += file.length;
        }
        size
    }

    pub fn add_peer(&mut self, peer: &mut Peer) -> io::Result<()> {
        peer.set_torrent(&self)?;
        self.peers.insert(peer.id);
        Ok(())
    }

    pub fn remove_peer(&mut self, peer: &mut Peer) {
        self.peers.remove(&peer.id);
    }
}
<|endoftext|>


<commit_before>use std::rc::*;

use super::address::*;
use super::treenode::*;
use super::basictree::*;

///
/// Represents which of the root's references have changed
///
pub enum TreeChangeType {
    /// The node's child reference has been replaced
    Child,

    /// The node's sibling reference has been replaced
    Sibling
}

///
/// A change represents an alteration to the tree
///
pub struct TreeChange {
    /// The address of the node matching the root of the change
    ///
    /// The address is relative to an imaginary node that is the parent of the real root node. This makes it possible to
    /// replace the entire tree by setting this to `TreeAddress::Here` and the change_type to `TreeChangeType::Child`.
    /// The real root node can be addressed at `TreeAddress::ChildAtIndex(0)`
    root: TreeAddress,

    /// Which of the root's references have changed
    change_type: TreeChangeType,

    /// The tree that should replace the changed reference. The last node in this tree (depth-first) will be given the same sibling as the last node in the replacement range
    replacement_tree: Option<TreeRef>
}

impl TreeChange {
    ///
    /// Creates a new tree change
    ///
    pub fn new<TAddress: ToTreeAddress, TNode: ToTreeNode>(root: &TAddress, change_type: TreeChangeType, replacement_tree: Option<&TNode>) -> TreeChange {
        TreeChange { root: root.to_tree_address(), change_type: change_type, replacement_tree: replacement_tree.map(|x| x.to_tree_node()) }
    }

    ///
    /// Performs a replacement on a basic tree node
    ///
    #[inline]
    fn perform_replacement(node: &BasicTree, change_type: &TreeChangeType, replacement_tree: &Option<TreeRef>) {
        match *change_type {
            TreeChangeType::Child => {
                match *replacement_tree {
                    Some(ref new_child)     => node.set_child_ref(new_child.clone()),
                    None                    => node.clear_child()
                }
            },

            TreeChangeType::Sibling => {
                match *replacement_tree {
                    Some(ref new_sibling)   => node.set_sibling_ref(new_sibling.clone()),
                    None                    => node.clear_child()
                }
            }
        }
    }

    ///
    /// Applies a change to a tree, generating a new tree
    ///
    fn perform_apply(original_tree: &TreeRef, address: &TreeAddress, change_type: &TreeChangeType, replacement_tree: &Option<TreeRef>) -> TreeRef {
        match *address {
            TreeAddress::Here => {
                let new_node = BasicTree::from(original_tree);
                TreeChange::perform_replacement(&new_node, change_type, replacement_tree);

                Rc::new(new_node)
            },

            TreeAddress::ChildAtIndex(child_index, ref child_address) => {
                // Copy the siblings into a stack
                let mut siblings    = vec![];
                let mut current     = original_tree.get_child_ref();

                for _ in 0..child_index {
                    siblings.push(current.clone().unwrap_or_else(|| Rc::new(BasicTree::new("", ()))));

                    current = current.and_then(|x| x.get_sibling_ref());
                }

                // Replace the child matching this item
                let child_tree  = current.clone().unwrap_or_else(|| Rc::new(BasicTree::new("", ())));
                let new_child   = TreeChange::perform_apply(&child_tree, &*child_address, change_type, replacement_tree);

                current = Some(new_child);

                // Pop siblings to generate the new child item
                while let Some(sibling) = siblings.pop() {
                    match current {
                        Some(next_sibling)  => current = Some(Rc::new(BasicTree::from_with_sibling(sibling, next_sibling))),
                        None                => current = Some(sibling),
                    }
                }

                Rc::new(BasicTree::from_with_child(original_tree, current.unwrap()))
            },

            TreeAddress::ChildWithTag(ref child_tag, ref child_address) => {
                // Copy the siblings into a stack
                let mut siblings    = vec![];
                let mut current     = original_tree.get_child_ref();

                while current.is_some() {
                    if let Some(ref child_ref) = current {
                        // Always true, but it isn't clear how to make while let work here when we need to update the pointer (because it's always borrowed in the loop block)
                        if child_ref.get_tag() == child_tag {
                            break;
                        }

                        siblings.push(child_ref.clone());
                    }

                    current = current.and_then(|x| x.get_sibling_ref());
                }

                // Replace the child matching this item
                let child_tree  = current.clone().unwrap_or_else(|| Rc::new(BasicTree::new("", ())));
                let new_child   = TreeChange::perform_apply(&child_tree, &*child_address, change_type, replacement_tree);

                current = Some(new_child);

                // Pop siblings to generate the new child item
                while let Some(sibling) = siblings.pop() {
                    match current {
                        Some(next_sibling)  => current = Some(Rc::new(BasicTree::from_with_sibling(sibling, next_sibling))),
                        None                => current = Some(sibling),
                    }
                }

                Rc::new(BasicTree::from_with_child(original_tree, current.unwrap()))
            }
        }
    }

    ///
    /// Creates a new tree from an old tree with this change added to it
    ///
    #[inline]
    pub fn apply(&self, original_tree: &TreeRef) -> TreeRef {
        // For simplicity, we actualise the 'imaginary' root
        let imaginary_root      = tree!("", original_tree).to_tree_node();
        let new_imaginary_root  = TreeChange::perform_apply(&imaginary_root, &self.root, &self.change_type, &self.replacement_tree);

        // Result is the child of the imaginary node (or an empty node if the entire tree is deleted)
        new_imaginary_root.get_child_ref().unwrap_or_else(|| "".to_tree_node())
    }

    ///
    /// Returns whether or not this change covers the specified address (or false if this cannot be determined)
    ///
    pub fn applies_to(&self, address: &TreeAddress) -> Option<bool> {
        match self.change_type {
            // If the child has changed, then anything that's a child of the root address is changed
            TreeChangeType::Child => self.root.to_tree_address().is_parent_of(address).map(|is_parent| { is_parent && self.root != *address }),

            // If the sibling has changed, then it's the parent address that's changed
            TreeChangeType::Sibling => self.root.parent().is_parent_of(address)
        }
    }
}

#[cfg(test)]
mod change_tests {
    use super::super::super::tree::*;

    #[test]
    fn can_apply_simple_change_tagged() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&("test", "one"), TreeChangeType::Sibling, Some(&("replaced", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(changed_tree.get_child_ref_at("one").unwrap().get_value().to_int(0) == 1);
        assert!(changed_tree.get_child_ref_at("replaced").unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at("replaced").unwrap().get_sibling_ref().is_none());
        assert!(changed_tree.get_child_ref_at("two").is_none());
        assert!(changed_tree.get_child_ref_at("three").is_none());
    }

    #[test]
    fn can_apply_simple_change_indexed() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&(0, 1), TreeChangeType::Sibling, Some(&("replaced", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(changed_tree.get_child_ref_at(0).unwrap().get_value().to_int(0) == 1);
        assert!(changed_tree.get_child_ref_at(1).unwrap().get_value().to_int(0) == 2);
        assert!(changed_tree.get_child_ref_at(2).unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at(2).unwrap().get_sibling_ref().is_none());
        assert!(changed_tree.get_child_ref_at(3).is_none());
    }

    #[test]
    fn can_apply_child_change_tagged() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&("test", "two").to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(!changed_tree.get_child_ref_at(("two", "new_child").to_tree_address()).is_none());

        assert!(changed_tree.get_child_ref_at(("two", "new_child").to_tree_address()).unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at(("two", "new_child").to_tree_address()).unwrap().get_sibling_ref().is_none());

        assert!(!changed_tree.get_child_ref_at("one").is_none());
        assert!(!changed_tree.get_child_ref_at("two").is_none());
        assert!(!changed_tree.get_child_ref_at("three").is_none());
        assert!(changed_tree.get_child_ref_at("two").unwrap().get_value().to_int(0) == 2);
    }

    #[test]
    fn can_apply_child_change_indexed() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(!changed_tree.get_child_ref_at((1, 0).to_tree_address()).is_none());

        assert!(changed_tree.get_child_ref_at((1, 0).to_tree_address()).unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at((1, 0).to_tree_address()).unwrap().get_sibling_ref().is_none());

        assert!(!changed_tree.get_child_ref_at(0).is_none());
        assert!(!changed_tree.get_child_ref_at(1).is_none());
        assert!(!changed_tree.get_child_ref_at(2).is_none());
        assert!(changed_tree.get_child_ref_at(1).unwrap().get_value().to_int(0) == 2);
    }

    #[test]
    fn child_change_applies_to_child_tree() {
        let change = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));

        assert!(change.applies_to(&(0, (1, 2)).to_tree_address()).unwrap());
    }

    #[test]
    fn child_change_does_not_apply_to_other_tree() {
        let change = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));

        assert!(!change.applies_to(&(0, (2, 2)).to_tree_address()).unwrap());
    }

    #[test]
    fn child_change_does_not_apply_to_parent() {
        let change = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));

        assert!(!change.applies_to(&(0, 1).to_tree_address()).unwrap());
    }

    #[test]
    fn sibling_change_applies_to_parent() {
        let change = TreeChange::new(&(0, (1, 2)).to_tree_address(), TreeChangeType::Sibling, Some(&("new_child", 4)));

        assert!(change.applies_to(&(0, 1).to_tree_address()).unwrap());
    }

    #[test]
    fn sibling_change_does_not_apply_to_different_parent() {
        let change = TreeChange::new(&(0, (1, 2)).to_tree_address(), TreeChangeType::Sibling, Some(&("new_child", 4)));

        assert!(!change.applies_to(&(0, 2).to_tree_address()).unwrap());
    }

    #[test]
    fn sibling_change_does_not_apply_to_grandparent() {
        let change = TreeChange::new(&(0, (1, 2)).to_tree_address(), TreeChangeType::Sibling, Some(&("new_child", 4)));

        assert!(!change.applies_to(&0.to_tree_address()).unwrap());
    }
}
<commit_msg>Change applies_to so that it's relative to the 'real' tree rather than having to incorporate the 'imaginary root' every time.<commit_after>use std::rc::*;

use super::address::*;
use super::treenode::*;
use super::basictree::*;

///
/// Represents which of the root's references have changed
///
pub enum TreeChangeType {
    /// The node's child reference has been replaced
    Child,

    /// The node's sibling reference has been replaced
    Sibling
}

///
/// A change represents an alteration to the tree
///
pub struct TreeChange {
    /// The address of the node matching the root of the change
    ///
    /// The address is relative to an imaginary node that is the parent of the real root node. This makes it possible to
    /// replace the entire tree by setting this to `TreeAddress::Here` and the change_type to `TreeChangeType::Child`.
    /// The real root node can be addressed at `TreeAddress::ChildAtIndex(0)`
    root: TreeAddress,

    /// Which of the root's references have changed
    change_type: TreeChangeType,

    /// The tree that should replace the changed reference. The last node in this tree (depth-first) will be given the same sibling as the last node in the replacement range
    replacement_tree: Option<TreeRef>
}

impl TreeChange {
    ///
    /// Creates a new tree change
    ///
    pub fn new<TAddress: ToTreeAddress, TNode: ToTreeNode>(root: &TAddress, change_type: TreeChangeType, replacement_tree: Option<&TNode>) -> TreeChange {
        TreeChange { root: root.to_tree_address(), change_type: change_type, replacement_tree: replacement_tree.map(|x| x.to_tree_node()) }
    }

    ///
    /// Performs a replacement on a basic tree node
    ///
    #[inline]
    fn perform_replacement(node: &BasicTree, change_type: &TreeChangeType, replacement_tree: &Option<TreeRef>) {
        match *change_type {
            TreeChangeType::Child => {
                match *replacement_tree {
                    Some(ref new_child)     => node.set_child_ref(new_child.clone()),
                    None                    => node.clear_child()
                }
            },

            TreeChangeType::Sibling => {
                match *replacement_tree {
                    Some(ref new_sibling)   => node.set_sibling_ref(new_sibling.clone()),
                    None                    => node.clear_child()
                }
            }
        }
    }

    ///
    /// Applies a change to a tree, generating a new tree
    ///
    fn perform_apply(original_tree: &TreeRef, address: &TreeAddress, change_type: &TreeChangeType, replacement_tree: &Option<TreeRef>) -> TreeRef {
        match *address {
            TreeAddress::Here => {
                let new_node = BasicTree::from(original_tree);
                TreeChange::perform_replacement(&new_node, change_type, replacement_tree);

                Rc::new(new_node)
            },

            TreeAddress::ChildAtIndex(child_index, ref child_address) => {
                // Copy the siblings into a stack
                let mut siblings    = vec![];
                let mut current     = original_tree.get_child_ref();

                for _ in 0..child_index {
                    siblings.push(current.clone().unwrap_or_else(|| Rc::new(BasicTree::new("", ()))));

                    current = current.and_then(|x| x.get_sibling_ref());
                }

                // Replace the child matching this item
                let child_tree  = current.clone().unwrap_or_else(|| Rc::new(BasicTree::new("", ())));
                let new_child   = TreeChange::perform_apply(&child_tree, &*child_address, change_type, replacement_tree);

                current = Some(new_child);

                // Pop siblings to generate the new child item
                while let Some(sibling) = siblings.pop() {
                    match current {
                        Some(next_sibling)  => current = Some(Rc::new(BasicTree::from_with_sibling(sibling, next_sibling))),
                        None                => current = Some(sibling),
                    }
                }

                Rc::new(BasicTree::from_with_child(original_tree, current.unwrap()))
            },

            TreeAddress::ChildWithTag(ref child_tag, ref child_address) => {
                // Copy the siblings into a stack
                let mut siblings    = vec![];
                let mut current     = original_tree.get_child_ref();

                while current.is_some() {
                    if let Some(ref child_ref) = current {
                        // Always true, but it isn't clear how to make while let work here when we need to update the pointer (because it's always borrowed in the loop block)
                        if child_ref.get_tag() == child_tag {
                            break;
                        }

                        siblings.push(child_ref.clone());
                    }

                    current = current.and_then(|x| x.get_sibling_ref());
                }

                // Replace the child matching this item
                let child_tree  = current.clone().unwrap_or_else(|| Rc::new(BasicTree::new("", ())));
                let new_child   = TreeChange::perform_apply(&child_tree, &*child_address, change_type, replacement_tree);

                current = Some(new_child);

                // Pop siblings to generate the new child item
                while let Some(sibling) = siblings.pop() {
                    match current {
                        Some(next_sibling)  => current = Some(Rc::new(BasicTree::from_with_sibling(sibling, next_sibling))),
                        None                => current = Some(sibling),
                    }
                }

                Rc::new(BasicTree::from_with_child(original_tree, current.unwrap()))
            }
        }
    }

    ///
    /// Creates a new tree from an old tree with this change added to it
    ///
    #[inline]
    pub fn apply(&self, original_tree: &TreeRef) -> TreeRef {
        // For simplicity, we actualise the 'imaginary' root
        let imaginary_root      = tree!("", original_tree).to_tree_node();
        let new_imaginary_root  = TreeChange::perform_apply(&imaginary_root, &self.root, &self.change_type, &self.replacement_tree);

        // Result is the child of the imaginary node (or an empty node if the entire tree is deleted)
        new_imaginary_root.get_child_ref().unwrap_or_else(|| "".to_tree_node())
    }

    ///
    /// Converts the root address to one that's relative to the tree being changed
    ///
    /// The root address starts at an imaginary 'true' root (this makes it possible to specify a change that replaces the entire tree)
    ///
    #[inline]
    fn address_relative_to_tree_root(&self) -> TreeAddress {
        match self.root {
            TreeAddress::Here                           => TreeAddress::Here,
            TreeAddress::ChildAtIndex(0, ref address)   => (**address).clone(),
            _                                           => TreeAddress::Here
        }
    }

    ///
    /// Returns whether or not this change covers the specified address (or false if this cannot be determined)
    ///
    pub fn applies_to(&self, address: &TreeAddress) -> Option<bool> {
        let relative_root = self.address_relative_to_tree_root();

        match self.change_type {
            // If the child has changed, then anything that's a child of the root address is changed
            TreeChangeType::Child => relative_root.is_parent_of(address).map(|is_parent| { is_parent && relative_root != *address }),

            // If the sibling has changed, then it's the parent address that's changed
            TreeChangeType::Sibling => relative_root.parent().is_parent_of(address)
        }
    }
}

#[cfg(test)]
mod change_tests {
    use super::super::super::tree::*;

    #[test]
    fn can_apply_simple_change_tagged() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&("test", "one"), TreeChangeType::Sibling, Some(&("replaced", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(changed_tree.get_child_ref_at("one").unwrap().get_value().to_int(0) == 1);
        assert!(changed_tree.get_child_ref_at("replaced").unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at("replaced").unwrap().get_sibling_ref().is_none());
        assert!(changed_tree.get_child_ref_at("two").is_none());
        assert!(changed_tree.get_child_ref_at("three").is_none());
    }

    #[test]
    fn can_apply_simple_change_indexed() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&(0, 1), TreeChangeType::Sibling, Some(&("replaced", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(changed_tree.get_child_ref_at(0).unwrap().get_value().to_int(0) == 1);
        assert!(changed_tree.get_child_ref_at(1).unwrap().get_value().to_int(0) == 2);
        assert!(changed_tree.get_child_ref_at(2).unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at(2).unwrap().get_sibling_ref().is_none());
        assert!(changed_tree.get_child_ref_at(3).is_none());
    }

    #[test]
    fn can_apply_child_change_tagged() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&("test", "two").to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(!changed_tree.get_child_ref_at(("two", "new_child").to_tree_address()).is_none());

        assert!(changed_tree.get_child_ref_at(("two", "new_child").to_tree_address()).unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at(("two", "new_child").to_tree_address()).unwrap().get_sibling_ref().is_none());

        assert!(!changed_tree.get_child_ref_at("one").is_none());
        assert!(!changed_tree.get_child_ref_at("two").is_none());
        assert!(!changed_tree.get_child_ref_at("three").is_none());
        assert!(changed_tree.get_child_ref_at("two").unwrap().get_value().to_int(0) == 2);
    }

    #[test]
    fn can_apply_child_change_indexed() {
        let initial_tree    = tree!("test", ("one", 1), ("two", 2), ("three", 3));
        let change_two      = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));
        let changed_tree    = change_two.apply(&initial_tree);

        assert!(!changed_tree.get_child_ref_at((1, 0).to_tree_address()).is_none());

        assert!(changed_tree.get_child_ref_at((1, 0).to_tree_address()).unwrap().get_value().to_int(0) == 4);
        assert!(changed_tree.get_child_ref_at((1, 0).to_tree_address()).unwrap().get_sibling_ref().is_none());

        assert!(!changed_tree.get_child_ref_at(0).is_none());
        assert!(!changed_tree.get_child_ref_at(1).is_none());
        assert!(!changed_tree.get_child_ref_at(2).is_none());
        assert!(changed_tree.get_child_ref_at(1).unwrap().get_value().to_int(0) == 2);
    }

    #[test]
    fn true_root_applies_to_everything() {
        // The child of address 0 represents the entire tree
        let change = TreeChange::new(&0.to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));

        // Note that the tree change applies to (0, 1) but it's relative to an imaginary root
        assert!(change.applies_to(&().to_tree_address()).unwrap());
        assert!(change.applies_to(&(1).to_tree_address()).unwrap());
        assert!(change.applies_to(&(1, 2).to_tree_address()).unwrap());
    }

    #[test]
    fn child_change_applies_to_child_tree() {
        let change = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));

        // Note that the tree change applies to (0, 1) but it's relative to an imaginary root
        assert!(change.applies_to(&(1, 2).to_tree_address()).unwrap());
    }

    #[test]
    fn child_change_does_not_apply_to_parent() {
        let change = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));

        assert!(!change.applies_to(&1.to_tree_address()).unwrap());
    }

    #[test]
    fn child_change_does_not_apply_to_other_tree() {
        let change = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&("new_child", 4)));

        assert!(!change.applies_to(&(2, 2).to_tree_address()).unwrap());
    }

    #[test]
    fn sibling_change_applies_to_parent() {
        let change = TreeChange::new(&(0, (1, 2)).to_tree_address(), TreeChangeType::Sibling, Some(&("new_child", 4)));

        assert!(change.applies_to(&1.to_tree_address()).unwrap());
    }

    #[test]
    fn sibling_change_does_not_apply_to_different_parent() {
        let change = TreeChange::new(&(0, (1, 2)).to_tree_address(), TreeChangeType::Sibling, Some(&("new_child", 4)));

        assert!(!change.applies_to(&2.to_tree_address()).unwrap());
    }

    #[test]
    fn sibling_change_does_not_apply_to_grandparent() {
        let change = TreeChange::new(&(0, (1, (2, 3))).to_tree_address(), TreeChangeType::Sibling, Some(&("new_child", 4)));

        assert!(!change.applies_to(&1.to_tree_address()).unwrap());
    }
}
<|endoftext|>


<commit_before>//! Futures integration.

use crate::stream::Stream;
use crate::sync::Mutex;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{LocalWaker, Poll, Waker};

/// The storage of a stream future.
#[derive(Debug)]
struct StreamFutureStorage<T> {
    value: Option<T>,
    waker: Option<Waker>,
}

impl<T> StreamFutureStorage<T> {
    fn new() -> Arc<Mutex<Self>> {
        Arc::new(Mutex::new(StreamFutureStorage {
            value: None,
            waker: None,
        }))
    }
}

/// A future that waits for a stream value.
///
/// This is created by `Stream::next`.
#[derive(Debug)]
pub struct StreamFuture<T> {
    storage: Arc<Mutex<StreamFutureStorage<T>>>,
    stream: Stream<T>,
}

impl<T> StreamFuture<T> {
    /// Creates a future that returns the next value sent to this stream.
    pub(crate) fn new(stream: Stream<T>) -> Self
    where
        T: Clone + Send + 'static,
    {
        let storage = StreamFutureStorage::new();
        let st = storage.clone();
        stream.observe(move |val| {
            let mut storage = st.lock();
            storage.value = Some(val.into_owned());
            if let Some(waker) = &storage.waker {
                waker.wake();
            }
            false
        });
        StreamFuture { storage, stream }
    }

    /// Obtains the source stream.
    pub fn get_source(&self) -> &Stream<T> {
        &self.stream
    }
}

impl<T> Future for StreamFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
        let mut storage = self.storage.lock();
        if let Some(value) = storage.value.take() {
            Poll::Ready(value)
        } else {
            storage.waker = Some(lw.clone().into_waker());
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stream::Sink;
    use futures::executor::block_on;

    #[test]
    fn basic() {
        let sink = Sink::new();
        let future = StreamFuture::new(sink.stream());

        sink.send(42);
        sink.send(13);
        assert_eq!(block_on(future), 42);
    }
}
<commit_msg>implement Unpin<commit_after>//! Futures integration.

use crate::stream::Stream;
use crate::sync::Mutex;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{LocalWaker, Poll, Waker};

/// The storage of a stream future.
#[derive(Debug)]
struct StreamFutureStorage<T> {
    value: Option<T>,
    waker: Option<Waker>,
}

impl<T> StreamFutureStorage<T> {
    fn new() -> Arc<Mutex<Self>> {
        Arc::new(Mutex::new(StreamFutureStorage {
            value: None,
            waker: None,
        }))
    }
}

/// A future that waits for a stream value.
///
/// This is created by `Stream::next`.
#[derive(Debug)]
pub struct StreamFuture<T> {
    storage: Arc<Mutex<StreamFutureStorage<T>>>,
    stream: Stream<T>,
}

impl<T> StreamFuture<T> {
    /// Creates a future that returns the next value sent to this stream.
    pub(crate) fn new(stream: Stream<T>) -> Self
    where
        T: Clone + Send + 'static,
    {
        let storage = StreamFutureStorage::new();
        let st = storage.clone();
        stream.observe(move |val| {
            let mut storage = st.lock();
            storage.value = Some(val.into_owned());
            if let Some(waker) = &storage.waker {
                waker.wake();
            }
            false
        });
        StreamFuture { storage, stream }
    }

    /// Obtains the source stream.
    pub fn get_source(&self) -> &Stream<T> {
        &self.stream
    }
}

impl<T> Future for StreamFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
        let mut storage = self.storage.lock();
        if let Some(value) = storage.value.take() {
            Poll::Ready(value)
        } else {
            storage.waker = Some(lw.clone().into_waker());
            Poll::Pending
        }
    }
}

impl<T> Unpin for StreamFuture<T> {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stream::Sink;
    use futures::executor::block_on;

    #[test]
    fn basic() {
        let sink = Sink::new();
        let future = StreamFuture::new(sink.stream());

        sink.send(42);
        sink.send(13);
        assert_eq!(block_on(future), 42);
    }
}
<|endoftext|>