<commit_before>use std::cmp::min;
use std::ops::BitOr;
use std::num::FromPrimitive;
use std::fmt::{self, Debug, Formatter};
use std::iter::IteratorExt;
use std::collections::{BTreeMap, VecDeque};
use std::string::ToString;
use serialize::json::Json;

use self::TemplateElement::{RawString, Expression, HelperExpression,
                            HTMLExpression, HelperBlock, Comment};

#[derive(PartialEq, Clone)]
pub struct Template {
    pub elements: Vec<TemplateElement>
}

#[derive(PartialEq, Debug)]
enum ParserState {
    Text,
    HtmlExpression,
    Comment,
    HelperStart,
    HelperEnd,
    Expression,
    Invalid
}

#[derive(PartialEq, Clone)]
pub struct Helper {
    name: String,
    params: Vec<String>,
    hash: BTreeMap<String, Json>,
    template: Option<Template>,
    inverse: Option<Template>,
    block: bool
}

impl Helper {
    pub fn name(&self) -> &String {
        &self.name
    }

    pub fn params(&self) -> &Vec<String> {
        &self.params
    }

    pub fn hash(&self) -> &BTreeMap<String, Json> {
        &self.hash
    }

    pub fn template(&self) -> Option<&Template> {
        match self.template {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn inverse(&self) -> Option<&Template> {
        match self.inverse {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn is_block(&self) -> bool {
        self.block
    }
}

impl ToString for Helper {
    fn to_string(&self) -> String {

        let mut buf = String::new();

        if self.block {
            buf.push_str(format!("{{{{#{}", self.name).as_slice());
        } else {
            buf.push_str(format!("{{{{{}", self.name).as_slice());
        }

        for p in self.params.iter() {
            buf.push_str(format!(" {}", p).as_slice());
        }

        for k in self.hash.keys() {
            buf.push_str(format!(" {}={}", k, self.hash.get(k).unwrap()).as_slice());
        }

        buf.push_str("}}");

        if self.block {
            let tpl = self.template();
            if tpl.is_some() {
                buf.push_str(tpl.unwrap().to_string().as_slice());
            }
            let ivs = self.inverse();
            if ivs.is_some() {
                buf.push_str("{{else}}");
                buf.push_str(ivs.unwrap().to_string().as_slice());
            }
            buf.push_str(format!("{{{{/{}}}}}", self.name).as_slice());
        }
        buf
    }
}

#[derive(Debug, Copy)]
pub struct TemplateError;

impl Helper {
    fn parse(source: String, block: bool) -> Result<Helper, TemplateError> {
        let mut tokens = source.split(|&: c: char| c.is_whitespace())
            .filter(|s: &&str| !(*s).is_empty());

        let name = tokens.next();
        match name {
            Some(n) => {
                let mut params: Vec<String> = Vec::new();
                let mut hash: BTreeMap<String, Json> = BTreeMap::new();

                for t in tokens {
                    if t.contains_char('=') {
                        let kv = t.split('=').collect::<Vec<&str>>();
                        hash.insert(kv.get(0).unwrap().to_string(),
                                    kv.get(1).unwrap().parse::<Json>().unwrap());
                    } else {
                        params.push(t.to_string());
                    }
                }

                Ok(Helper{
                    name: n.to_string(),
                    params: params,
                    hash: hash,
                    template: Option::None,
                    inverse: Option::None,
                    block: block
                })
            },
            None =>
                Err(TemplateError)
        }
    }
}

#[derive(PartialEq, FromPrimitive)]
enum WhiteSpaceOmit {
    Left = 0x01,
    Right = 0x10,
    Both = 0x11,
    None = 0x00
}

impl BitOr<WhiteSpaceOmit> for WhiteSpaceOmit {
    type Output = WhiteSpaceOmit;

    fn bitor(self, right: WhiteSpaceOmit) -> WhiteSpaceOmit {
        FromPrimitive::from_u8((self as u8) | (right as u8)).unwrap()
    }
}

fn process_whitespace(buf: &String, wso: &mut WhiteSpaceOmit) -> String {
    let result = match *wso {
        WhiteSpaceOmit::Left => {
            buf.trim_left().to_string()
        },
        WhiteSpaceOmit::Right => {
            buf.trim_right().to_string()
        },
        WhiteSpaceOmit::Both => {
            buf.trim().to_string()
        },
        WhiteSpaceOmit::None => {
            buf.clone()
        }
    };
    *wso = WhiteSpaceOmit::None;
    result
}

impl Template {
    pub fn compile(source: String) -> Result<Template, TemplateError> {
        let mut helper_stack: VecDeque<Helper> = VecDeque::new();
        let mut template_stack: VecDeque<Template> = VecDeque::new();
        template_stack.push_front(Template{ elements: Vec::new() });

        let mut buffer: String = String::new();
        let mut state = ParserState::Text;

        let mut c:usize = 0;
        let mut ws_omitter = WhiteSpaceOmit::None;
        let source_len = source.chars().count();
        while c < source_len {
            let mut slice = source.slice_chars(c, min(c+3, source_len)).to_string();
            if slice == "{{~" {
                ws_omitter = ws_omitter | WhiteSpaceOmit::Right;
                // read another char and remove ~
                slice = source.slice_chars(c, min(c+4, source_len)).to_string();
                slice.remove(2);
                c += 1;
            }
            if slice == "~}}" {
                ws_omitter = ws_omitter | WhiteSpaceOmit::Left;
                c += 1;
                slice = source.slice_chars(c, min(c+3, source_len)).to_string();
            }
            state = match slice.as_slice() {
                "{{{" | "{{!" | "{{#" | "{{/" => {
                    c += 2;
                    if !buffer.is_empty() {
                        let mut t = template_stack.front_mut().unwrap();
                        let buf_clone = process_whitespace(&buffer, &mut ws_omitter);
                        t.elements.push(RawString(buf_clone));
                        buffer.clear();
                    }
                    match slice.as_slice() {
                        "{{{" => ParserState::HtmlExpression,
                        "{{!" => ParserState::Comment,
                        "{{#" => ParserState::HelperStart,
                        "{{/" => {
                            let t = template_stack.pop_front().unwrap();
                            let h = helper_stack.front_mut().unwrap();
                            if h.template.is_some() {
                                h.inverse = Some(t);
                            } else {
                                h.template = Some(t);
                            }
                            ParserState::HelperEnd
                        },
                        _ => ParserState::Invalid
                    }
                },
                "}}}" => {
                    c += 2;
                    let mut t = template_stack.front_mut().unwrap();
                    t.elements.push(HTMLExpression(
                        buffer.clone().trim_matches(' ').to_string()));
                    buffer.clear();
                    ParserState::Text
                },
                _ => {
                    match if slice.len() > 2 { slice.slice_chars(0, 2) } else { slice.as_slice() } {
                        "{{" => {
                            c += 1;
                            if !buffer.is_empty() {
                                let mut t = template_stack.front_mut().unwrap();
                                let buf_clone = process_whitespace(&buffer, &mut ws_omitter);
                                t.elements.push(RawString(buf_clone));
                                buffer.clear();
                            }
                            ParserState::Expression
                        },
                        "}}" => {
                            c += 1;
                            match state {
                                ParserState::Expression => {
                                    if !buffer.is_empty() {
                                        // {{else}} or {{^}} within a helper block
                                        if buffer.trim() == "else" || buffer.trim() == "^" {
                                            buffer.clear(); // drop else
                                            let t = template_stack.pop_front().unwrap();
                                            let h = helper_stack.front_mut().unwrap();
                                            h.template = Some(t);
                                            template_stack.push_front(Template{ elements: Vec::new() });
                                            ParserState::Text
                                        } else if buffer.contains_char(' ') {
                                            //inline helper
                                            match Helper::parse(buffer.clone(), false){
                                                Ok(helper) => {
                                                    let mut t = template_stack.front_mut().unwrap();
                                                    t.elements.push(HelperExpression(helper));
                                                    buffer.clear();
                                                    ParserState::Text
                                                },
                                                Err(_) => ParserState::Invalid
                                            }
                                        } else {
                                            let mut t = template_stack.front_mut().unwrap();
                                            t.elements.push(Expression(
                                                buffer.clone().trim_matches(' ').to_string()));
                                            buffer.clear();
                                            ParserState::Text
                                        }

                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                ParserState::Comment => {
                                    let mut t = template_stack.front_mut().unwrap();
                                    t.elements.push(Comment(buffer.clone()));
                                    buffer.clear();
                                    ParserState::Text
                                },
                                ParserState::HelperStart => {
                                    match Helper::parse(buffer.clone(), true) {
                                        Ok(helper) => {
                                            helper_stack.push_front(helper);
                                            template_stack.push_front(Template{ elements: Vec::new() });

                                            buffer.clear();
                                            ParserState::Text
                                        },
                                        Err(_) => ParserState::Invalid
                                    }
                                },
                                ParserState::HelperEnd => {
                                    let name = buffer.trim_matches(' ').to_string();
                                    if name == helper_stack.front().unwrap().name {
                                        let h = helper_stack.pop_front().unwrap();
                                        let mut t = template_stack.front_mut().unwrap();
                                        t.elements.push(HelperBlock(h));
                                        buffer.clear();
                                        ParserState::Text
                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                _ => ParserState::Invalid
                            }
                        },
                        _ => {
                            buffer.push(slice.char_at(0));
                            state
                        }
                    }
                }
            };
            if state == ParserState::Invalid {
                return Err(TemplateError);
            }
            c += 1;
        }

        if !buffer.is_empty() {
            let mut t = template_stack.front_mut().unwrap();
            let buf_clone = process_whitespace(&buffer, &mut ws_omitter);
            t.elements.push(TemplateElement::RawString(buf_clone));
        }

        if !helper_stack.is_empty() {
            return Err(TemplateError);
        }

        return Ok(template_stack.pop_front().unwrap());
    }
}

impl ToString for Template {
    fn to_string(&self) -> String {
        let mut buf = String::new();
        for v in self.elements.iter() {
            buf.push_str(v.to_string().as_slice());
        }
        buf
    }
}

#[derive(PartialEq, Clone)]
pub enum TemplateElement {
    RawString(String),
    Expression(String),
    HTMLExpression(String),
    HelperExpression(Helper),
    HelperBlock(Helper),
    Comment(String),
}

impl ToString for TemplateElement {
    fn to_string(&self) -> String {
        match *self {
            RawString(ref v) => {
                v.clone()
            },
            Expression(ref v) => {
                // {{ escape to {
                format!("{{{{{}}}}}", v)
            },
            HTMLExpression(ref v) => {
                format!("{{{{{{{}}}}}}}", v)
            },
            HelperExpression(ref helper) => {
                helper.to_string()
            }
            HelperBlock(ref helper) => {
                helper.to_string()
            }
            Comment(ref v) => {
                format!("{{!{}}}", v)
            }
        }
    }
}

impl Debug for TemplateElement {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        try!(writeln!(f, "{:?}", self.to_string()));
        Ok(())
    }
}

#[test]
fn test_parse_helper_start_tag() {
    let source = "if not name compare=1".to_string();
    let h = Helper::parse(source, true).ok().unwrap();

    assert_eq!(h.name, "if".to_string());
    assert_eq!(h.params, vec::<String>!["not".to_string(), "name".to_string()]);

    let key = "compare".to_string();
    let value = h.hash.get(&key).unwrap();
    assert_eq!(value.as_u64().unwrap(), 1u64);
}

#[test]
fn test_parse_template() {
    let source = "
{{title}} 你好
 {{{content}}}
{{#if date}}
good
{{else}}
bad
{{/if}}{{foo bar}}{{#unless true}}kitkat{{^}}lollipop{{/unless}}";
    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 9);
    assert_eq!((*t.elements.get(0).unwrap()).to_string(), "
".to_string());
    assert_eq!((*t.elements.get(1).unwrap()).to_string(),
               Expression("title".to_string()).to_string());

    assert_eq!((*t.elements.get(3).unwrap()).to_string(), "{{{content}}}".to_string());

    match *t.elements.get(5).unwrap() {
        HelperBlock(ref h) => {
            assert_eq!(h.name, "if".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(h.template.as_ref().unwrap().elements.len(), 1);
        },
        _ => {
            panic!("Helper expected here.");
        }
    };

    match *t.elements.get(7).unwrap() {
        HelperExpression(ref h) => {
            assert_eq!(h.name, "foo".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(*(h.params.get(0).unwrap()), "bar".to_string());
        },
        _ => {
            panic!("Helper expression here");
        }
    };

    match *t.elements.get(8).unwrap() {
        HelperBlock(ref h) => {
            assert_eq!(h.name, "unless".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(h.inverse.as_ref().unwrap().elements.len(), 1);
        },
        _ => {
            panic!("Helper expression here");
        }
    };

}

#[test]
fn test_helper_to_string() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good{{/ifequals}}".to_string();

    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 1);
    assert_eq!(t.elements.get(0).unwrap().to_string(), source);
}

#[test]
fn test_parse_error() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good";

    let t = Template::compile(source.to_string());

    assert!(t.is_err());
}

#[test]
fn test_white_space_omitter() {
    let source = "hello~     {{~world~}} \n  !{{~#if true}}else{{/if~}}".to_string();
    let t = Template::compile(source).ok().unwrap();

    assert_eq!(t.elements.len(), 4);

    assert_eq!(t.elements[0], RawString(String::from_str("hello~")));
    assert_eq!(t.elements[1], Expression(String::from_str("world")));
    assert_eq!(t.elements[2], RawString(String::from_str("!")));
}
<commit_msg>(fix) rustup<commit_after>use std::cmp::min;
use std::ops::BitOr;
use std::num::FromPrimitive;
use std::fmt::{self, Debug, Formatter};
use std::iter::IteratorExt;
use std::collections::{BTreeMap, VecDeque};
use std::string::ToString;
use serialize::json::Json;

use self::TemplateElement::{RawString, Expression, HelperExpression,
                            HTMLExpression, HelperBlock, Comment};

#[derive(PartialEq, Clone)]
pub struct Template {
    pub elements: Vec<TemplateElement>
}

#[derive(PartialEq, Debug)]
enum ParserState {
    Text,
    HtmlExpression,
    Comment,
    HelperStart,
    HelperEnd,
    Expression,
    Invalid
}

#[derive(PartialEq, Clone)]
pub struct Helper {
    name: String,
    params: Vec<String>,
    hash: BTreeMap<String, Json>,
    template: Option<Template>,
    inverse: Option<Template>,
    block: bool
}

impl Helper {
    pub fn name(&self) -> &String {
        &self.name
    }

    pub fn params(&self) -> &Vec<String> {
        &self.params
    }

    pub fn hash(&self) -> &BTreeMap<String, Json> {
        &self.hash
    }

    pub fn template(&self) -> Option<&Template> {
        match self.template {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn inverse(&self) -> Option<&Template> {
        match self.inverse {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn is_block(&self) -> bool {
        self.block
    }
}

impl ToString for Helper {
    fn to_string(&self) -> String {

        let mut buf = String::new();

        if self.block {
            buf.push_str(format!("{{{{#{}", self.name).as_slice());
        } else {
            buf.push_str(format!("{{{{{}", self.name).as_slice());
        }

        for p in self.params.iter() {
            buf.push_str(format!(" {}", p).as_slice());
        }

        for k in self.hash.keys() {
            buf.push_str(format!(" {}={}", k, self.hash.get(k).unwrap()).as_slice());
        }

        buf.push_str("}}");

        if self.block {
            let tpl = self.template();
            if tpl.is_some() {
                buf.push_str(tpl.unwrap().to_string().as_slice());
            }
            let ivs = self.inverse();
            if ivs.is_some() {
                buf.push_str("{{else}}");
                buf.push_str(ivs.unwrap().to_string().as_slice());
            }
            buf.push_str(format!("{{{{/{}}}}}", self.name).as_slice());
        }
        buf
    }
}

#[derive(Debug, Copy)]
pub struct TemplateError;

impl Helper {
    fn parse(source: String, block: bool) -> Result<Helper, TemplateError> {
        let mut tokens = source.split(|&: c: char| c.is_whitespace())
            .filter(|s: &&str| !(*s).is_empty());

        let name = tokens.next();
        match name {
            Some(n) => {
                let mut params: Vec<String> = Vec::new();
                let mut hash: BTreeMap<String, Json> = BTreeMap::new();

                for t in tokens {
                    if t.contains('=') {
                        let kv = t.split('=').collect::<Vec<&str>>();
                        hash.insert(kv.get(0).unwrap().to_string(),
                                    kv.get(1).unwrap().parse::<Json>().unwrap());
                    } else {
                        params.push(t.to_string());
                    }
                }

                Ok(Helper{
                    name: n.to_string(),
                    params: params,
                    hash: hash,
                    template: Option::None,
                    inverse: Option::None,
                    block: block
                })
            },
            None =>
                Err(TemplateError)
        }
    }
}

#[derive(PartialEq, FromPrimitive)]
enum WhiteSpaceOmit {
    Left = 0x01,
    Right = 0x10,
    Both = 0x11,
    None = 0x00
}

impl BitOr<WhiteSpaceOmit> for WhiteSpaceOmit {
    type Output = WhiteSpaceOmit;

    fn bitor(self, right: WhiteSpaceOmit) -> WhiteSpaceOmit {
        FromPrimitive::from_u8((self as u8) | (right as u8)).unwrap()
    }
}

fn process_whitespace(buf: &String, wso: &mut WhiteSpaceOmit) -> String {
    let result = match *wso {
        WhiteSpaceOmit::Left => {
            buf.trim_left().to_string()
        },
        WhiteSpaceOmit::Right => {
            buf.trim_right().to_string()
        },
        WhiteSpaceOmit::Both => {
            buf.trim().to_string()
        },
        WhiteSpaceOmit::None => {
            buf.clone()
        }
    };
    *wso = WhiteSpaceOmit::None;
    result
}

impl Template {
    pub fn compile(source: String) -> Result<Template, TemplateError> {
        let mut helper_stack: VecDeque<Helper> = VecDeque::new();
        let mut template_stack: VecDeque<Template> = VecDeque::new();
        template_stack.push_front(Template{ elements: Vec::new() });

        let mut buffer: String = String::new();
        let mut state = ParserState::Text;

        let mut c:usize = 0;
        let mut ws_omitter = WhiteSpaceOmit::None;
        let source_len = source.chars().count();
        while c < source_len {
            let mut slice = source.slice_chars(c, min(c+3, source_len)).to_string();
            if slice == "{{~" {
                ws_omitter = ws_omitter | WhiteSpaceOmit::Right;
                // read another char and remove ~
                slice = source.slice_chars(c, min(c+4, source_len)).to_string();
                slice.remove(2);
                c += 1;
            }
            if slice == "~}}" {
                ws_omitter = ws_omitter | WhiteSpaceOmit::Left;
                c += 1;
                slice = source.slice_chars(c, min(c+3, source_len)).to_string();
            }
            state = match slice.as_slice() {
                "{{{" | "{{!" | "{{#" | "{{/" => {
                    c += 2;
                    if !buffer.is_empty() {
                        let mut t = template_stack.front_mut().unwrap();
                        let buf_clone = process_whitespace(&buffer, &mut ws_omitter);
                        t.elements.push(RawString(buf_clone));
                        buffer.clear();
                    }
                    match slice.as_slice() {
                        "{{{" => ParserState::HtmlExpression,
                        "{{!" => ParserState::Comment,
                        "{{#" => ParserState::HelperStart,
                        "{{/" => {
                            let t = template_stack.pop_front().unwrap();
                            let h = helper_stack.front_mut().unwrap();
                            if h.template.is_some() {
                                h.inverse = Some(t);
                            } else {
                                h.template = Some(t);
                            }
                            ParserState::HelperEnd
                        },
                        _ => ParserState::Invalid
                    }
                },
                "}}}" => {
                    c += 2;
                    let mut t = template_stack.front_mut().unwrap();
                    t.elements.push(HTMLExpression(
                        buffer.clone().trim_matches(' ').to_string()));
                    buffer.clear();
                    ParserState::Text
                },
                _ => {
                    match if slice.len() > 2 { slice.slice_chars(0, 2) } else { slice.as_slice() } {
                        "{{" => {
                            c += 1;
                            if !buffer.is_empty() {
                                let mut t = template_stack.front_mut().unwrap();
                                let buf_clone = process_whitespace(&buffer, &mut ws_omitter);
                                t.elements.push(RawString(buf_clone));
                                buffer.clear();
                            }
                            ParserState::Expression
                        },
                        "}}" => {
                            c += 1;
                            match state {
                                ParserState::Expression => {
                                    if !buffer.is_empty() {
                                        // {{else}} or {{^}} within a helper block
                                        if buffer.trim() == "else" || buffer.trim() == "^" {
                                            buffer.clear(); // drop else
                                            let t = template_stack.pop_front().unwrap();
                                            let h = helper_stack.front_mut().unwrap();
                                            h.template = Some(t);
                                            template_stack.push_front(Template{ elements: Vec::new() });
                                            ParserState::Text
                                        } else if buffer.contains(' ') {
                                            //inline helper
                                            match Helper::parse(buffer.clone(), false){
                                                Ok(helper) => {
                                                    let mut t = template_stack.front_mut().unwrap();
                                                    t.elements.push(HelperExpression(helper));
                                                    buffer.clear();
                                                    ParserState::Text
                                                },
                                                Err(_) => ParserState::Invalid
                                            }
                                        } else {
                                            let mut t = template_stack.front_mut().unwrap();
                                            t.elements.push(Expression(
                                                buffer.clone().trim_matches(' ').to_string()));
                                            buffer.clear();
                                            ParserState::Text
                                        }

                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                ParserState::Comment => {
                                    let mut t = template_stack.front_mut().unwrap();
                                    t.elements.push(Comment(buffer.clone()));
                                    buffer.clear();
                                    ParserState::Text
                                },
                                ParserState::HelperStart => {
                                    match Helper::parse(buffer.clone(), true) {
                                        Ok(helper) => {
                                            helper_stack.push_front(helper);
                                            template_stack.push_front(Template{ elements: Vec::new() });

                                            buffer.clear();
                                            ParserState::Text
                                        },
                                        Err(_) => ParserState::Invalid
                                    }
                                },
                                ParserState::HelperEnd => {
                                    let name = buffer.trim_matches(' ').to_string();
                                    if name == helper_stack.front().unwrap().name {
                                        let h = helper_stack.pop_front().unwrap();
                                        let mut t = template_stack.front_mut().unwrap();
                                        t.elements.push(HelperBlock(h));
                                        buffer.clear();
                                        ParserState::Text
                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                _ => ParserState::Invalid
                            }
                        },
                        _ => {
                            buffer.push(slice.char_at(0));
                            state
                        }
                    }
                }
            };
            if state == ParserState::Invalid {
                return Err(TemplateError);
            }
            c += 1;
        }

        if !buffer.is_empty() {
            let mut t = template_stack.front_mut().unwrap();
            let buf_clone = process_whitespace(&buffer, &mut ws_omitter);
            t.elements.push(TemplateElement::RawString(buf_clone));
        }

        if !helper_stack.is_empty() {
            return Err(TemplateError);
        }

        return Ok(template_stack.pop_front().unwrap());
    }
}

impl ToString for Template {
    fn to_string(&self) -> String {
        let mut buf = String::new();
        for v in self.elements.iter() {
            buf.push_str(v.to_string().as_slice());
        }
        buf
    }
}

#[derive(PartialEq, Clone)]
pub enum TemplateElement {
    RawString(String),
    Expression(String),
    HTMLExpression(String),
    HelperExpression(Helper),
    HelperBlock(Helper),
    Comment(String),
}

impl ToString for TemplateElement {
    fn to_string(&self) -> String {
        match *self {
            RawString(ref v) => {
                v.clone()
            },
            Expression(ref v) => {
                // {{ escape to {
                format!("{{{{{}}}}}", v)
            },
            HTMLExpression(ref v) => {
                format!("{{{{{{{}}}}}}}", v)
            },
            HelperExpression(ref helper) => {
                helper.to_string()
            }
            HelperBlock(ref helper) => {
                helper.to_string()
            }
            Comment(ref v) => {
                format!("{{!{}}}", v)
            }
        }
    }
}

impl Debug for TemplateElement {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        try!(writeln!(f, "{:?}", self.to_string()));
        Ok(())
    }
}

#[test]
fn test_parse_helper_start_tag() {
    let source = "if not name compare=1".to_string();
    let h = Helper::parse(source, true).ok().unwrap();

    assert_eq!(h.name, "if".to_string());
    assert_eq!(h.params, vec::<String>!["not".to_string(), "name".to_string()]);

    let key = "compare".to_string();
    let value = h.hash.get(&key).unwrap();
    assert_eq!(value.as_u64().unwrap(), 1u64);
}

#[test]
fn test_parse_template() {
    let source = "
{{title}} 你好
 {{{content}}}
{{#if date}}
good
{{else}}
bad
{{/if}}{{foo bar}}{{#unless true}}kitkat{{^}}lollipop{{/unless}}";
    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 9);
    assert_eq!((*t.elements.get(0).unwrap()).to_string(), "
".to_string());
    assert_eq!((*t.elements.get(1).unwrap()).to_string(),
               Expression("title".to_string()).to_string());

    assert_eq!((*t.elements.get(3).unwrap()).to_string(), "{{{content}}}".to_string());

    match *t.elements.get(5).unwrap() {
        HelperBlock(ref h) => {
            assert_eq!(h.name, "if".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(h.template.as_ref().unwrap().elements.len(), 1);
        },
        _ => {
            panic!("Helper expected here.");
        }
    };

    match *t.elements.get(7).unwrap() {
        HelperExpression(ref h) => {
            assert_eq!(h.name, "foo".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(*(h.params.get(0).unwrap()), "bar".to_string());
        },
        _ => {
            panic!("Helper expression here");
        }
    };

    match *t.elements.get(8).unwrap() {
        HelperBlock(ref h) => {
            assert_eq!(h.name, "unless".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(h.inverse.as_ref().unwrap().elements.len(), 1);
        },
        _ => {
            panic!("Helper expression here");
        }
    };

}

#[test]
fn test_helper_to_string() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good{{/ifequals}}".to_string();

    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 1);
    assert_eq!(t.elements.get(0).unwrap().to_string(), source);
}

#[test]
fn test_parse_error() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good";

    let t = Template::compile(source.to_string());

    assert!(t.is_err());
}

#[test]
fn test_white_space_omitter() {
    let source = "hello~     {{~world~}} \n  !{{~#if true}}else{{/if~}}".to_string();
    let t = Template::compile(source).ok().unwrap();

    assert_eq!(t.elements.len(), 4);

    assert_eq!(t.elements[0], RawString(String::from_str("hello~")));
    assert_eq!(t.elements[1], Expression(String::from_str("world")));
    assert_eq!(t.elements[2], RawString(String::from_str("!")));
}
<|endoftext|>


<commit_before>use Error::{self, ArgumentError, OperationError};
use Result;

use bson::{self, Bson, oid};
use chrono::{DateTime, UTC};

use coll::options::FindOptions;
use connstring::{self, ConnectionString, Host};
use cursor::Cursor;
use pool::{ConnectionPool, PooledStream};
use wire_protocol::flags::OpQueryFlags;

use std::collections::{BTreeMap, HashMap};
use std::net::TcpStream;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicIsize, Ordering};
use std::thread;

const DEFAULT_HEARTBEAT_FREQUENCY_MS: u32 = 10000;
const DEFAULT_RETRY_FREQUENCY_MS: u32 = 1000;
const DEFAULT_MAX_BSON_OBJECT_SIZE: i64 = 16 * 1024 * 1024;
const DEFAULT_MAX_MESSAGE_SIZE_BYTES: i64 = 48000000;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TopologyType {
    Single,
    ReplicaSetNoPrimary,
    ReplicaSetWithPrimary,
    Sharded,
    Unknown,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ServerType {
    Standalone,
    Mongos,
    RSPrimary,
    RSSecondary,
    RSArbiter,
    RSOther,
    RSGhost,
    Unknown,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct IsMasterResult {
    is_master: bool,
    max_bson_object_size: i64,
    max_message_size_bytes: i64,
    local_time: DateTime<UTC>,
    min_wire_version: i64,
    max_wire_version: i64,

    // Shards
    msg: String,

    // RS
    is_replica_set: bool,
    is_secondary: bool,
    me: Option<Host>,
    hosts: Vec<Host>,
    passives: Vec<Host>,
    arbiters: Vec<Host>,
    arbiter_only: bool,
    tags: BTreeMap<String, String>,
    set_name: String,
    election_id: Option<oid::ObjectId>,
    primary: Option<Host>,
    hidden: bool,
}

#[derive(Clone)]
pub struct Topology {
    config: ConnectionString,
    description: TopologyDescription,
    servers: Arc<HashMap<Host, RwLock<Server>>>,
    compatible: bool,
    compat_error: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TopologyDescription {
    ttype: TopologyType,
    set_name: String,
    heartbeat_frequency_ms: u32,
    max_election_id: Option<oid::ObjectId>,
}

#[derive(Clone)]
pub struct Server {
    pub host: Host,
    pool: ConnectionPool,
    description: Arc<RwLock<ServerDescription>>,
}

#[derive(Clone, Debug)]
pub struct ServerDescription {
    stype: ServerType,
    err: Arc<Option<Error>>,
    round_trip_time: Option<i64>,
    min_wire_version: i64,
    max_wire_version: i64,
    me: Option<Host>,
    hosts: Vec<Host>,
    passives: Vec<Host>,
    arbiters: Vec<Host>,
    tags: BTreeMap<String, String>,
    set_name: String,
    election_id: Option<oid::ObjectId>,
    primary: Option<Host>,
}

struct Monitor {
    host: Host,
    description: Arc<RwLock<ServerDescription>>,
    socket: TcpStream,
}

impl IsMasterResult {
    pub fn new(doc: bson::Document) -> Result<IsMasterResult> {
        let is_master = match doc.get("ismaster") {
            Some(&Bson::Boolean(ref b)) => *b,
            _ => return Err(ArgumentError("result does not contain 'ismaster'.".to_owned())),
        };

        let local_time = match doc.get("localTime") {
            Some(&Bson::UtcDatetime(ref datetime)) => datetime.clone(),
            _ => return Err(ArgumentError("result does not contain 'localTime'.".to_owned())),
        };

        let min_version = match doc.get("minWireVersion") {
            Some(&Bson::I64(ref v)) => *v,
            _ => return Err(ArgumentError("result does not contain 'minWireVersion'.".to_owned())),
        };

        let max_version = match doc.get("maxWireVersion") {
            Some(&Bson::I64(ref v)) => *v,
            _ => return Err(ArgumentError("result does not contain 'maxWireVersion'.".to_owned())),
        };

        let mut result = IsMasterResult {
            is_master: is_master,
            max_bson_object_size: DEFAULT_MAX_BSON_OBJECT_SIZE,
            max_message_size_bytes: DEFAULT_MAX_MESSAGE_SIZE_BYTES,
            local_time: local_time,
            min_wire_version: min_version,
            max_wire_version: max_version,
            msg: String::new(),
            is_secondary: false,
            is_replica_set: false,
            me: None,
            hosts: Vec::new(),
            passives: Vec::new(),
            arbiters: Vec::new(),
            arbiter_only: false,
            tags: BTreeMap::new(),
            set_name: String::new(),
            election_id: None,
            primary: None,
            hidden: false,
        };

        if let Some(&Bson::String(ref s)) = doc.get("msg") {
            result.msg = s.to_owned();
        }

        if let Some(&Bson::Boolean(ref b)) = doc.get("secondary") {
            result.is_secondary = *b;
        }

        if let Some(&Bson::Boolean(ref b)) = doc.get("isreplicaset") {
            result.is_replica_set = *b;
        }

        if let Some(&Bson::String(ref s)) = doc.get("setName") {
            result.set_name = s.to_owned();
        }

        if let Some(&Bson::String(ref s)) = doc.get("me") {
            result.me = Some(try!(connstring::parse_host(s)));
        }

        if let Some(&Bson::Array(ref arr)) = doc.get("hosts") {
            result.hosts = arr.iter().filter_map(|bson| match bson {
                &Bson::String(ref s) => connstring::parse_host(s).ok(),
                _ => None,
            }).collect();
        }

        if let Some(&Bson::Array(ref arr)) = doc.get("passives") {
            result.passives = arr.iter().filter_map(|bson| match bson {
                &Bson::String(ref s) => connstring::parse_host(s).ok(),
                _ => None,
            }).collect();
        }

        if let Some(&Bson::Array(ref arr)) = doc.get("arbiters") {
            result.passives = arr.iter().filter_map(|bson| match bson {
                &Bson::String(ref s) => connstring::parse_host(s).ok(),
                _ => None,
            }).collect();
        }

        if let Some(&Bson::String(ref s)) = doc.get("primary") {
            result.primary = Some(try!(connstring::parse_host(s)));
        }

        if let Some(&Bson::Boolean(ref b)) = doc.get("arbiterOnly") {
            result.arbiter_only = *b;
        }

        if let Some(&Bson::Boolean(ref h)) = doc.get("hidden") {
            result.hidden = *h;
        }

        if let Some(&Bson::Document(ref doc)) = doc.get("tags") {
            for (k, v) in doc.into_iter() {
                if let &Bson::String(ref tag) = v {
                    result.tags.insert(k.to_owned(), tag.to_owned());
                }
            }
        }

        if let Some(&Bson::ObjectId(ref id)) = doc.get("electionId") {
            result.election_id = Some(id.clone());
        }

        Ok(result)
    }
}

impl TopologyDescription {
    pub fn new() -> TopologyDescription {
        TopologyDescription {
            ttype: TopologyType::Unknown,
            set_name: String::new(),
            heartbeat_frequency_ms: DEFAULT_HEARTBEAT_FREQUENCY_MS,
            max_election_id: None,
        }
    }
}

impl Topology {
    pub fn new(req_id: Arc<AtomicIsize>, config: ConnectionString,
               description: Option<TopologyDescription>) -> Result<Topology> {

        let options = description.unwrap_or(TopologyDescription::new());

        if config.hosts.len() > 1 && options.ttype == TopologyType::Single {
            return Err(ArgumentError(
                "TopologyType::Single cannot be used with multiple seeds.".to_owned()));
        }

        if !options.set_name.is_empty() && options.ttype != TopologyType::ReplicaSetNoPrimary {
            return Err(ArgumentError(
                "TopologyType must be ReplicaSetNoPrimary if set_name is provided.".to_owned()));
        }

        // TODO: Determine driver's wire compatibility, and check overlap with
        // all servers in topology.

        let mut servers = HashMap::new();
        for host in config.hosts.iter() {
            let server = Server::new(req_id.clone(), host.clone());
            servers.insert(host.clone(), RwLock::new(server));
        }

        Ok(Topology {
            config: config,
            description: options,
            servers: Arc::new(servers),
            compatible: true,
            compat_error: String::new(),
        })
    }

    pub fn acquire_stream(&self) -> Result<PooledStream> {
        for (_, server) in self.servers.iter() {
            let read_server = try!(server.read());
            return read_server.acquire_stream()
        }
        Err(OperationError("No servers found in configuration.".to_owned()))
    }
}

impl ServerDescription {
    pub fn new() -> ServerDescription {
        ServerDescription {
            stype: ServerType::Unknown,
            err: Arc::new(None),
            round_trip_time: None,
            min_wire_version: 0,
            max_wire_version: 0,
            me: None,
            hosts: Vec::new(),
            passives: Vec::new(),
            arbiters: Vec::new(),
            tags: BTreeMap::new(),
            set_name: String::new(),
            election_id: None,
            primary: None,
        }
    }

    pub fn update(&mut self, ismaster: IsMasterResult) {
        self.min_wire_version = ismaster.min_wire_version;
        self.max_wire_version = ismaster.max_wire_version;
        self.me = ismaster.me;
        self.hosts = ismaster.hosts;
        self.passives = ismaster.passives;
        self.arbiters = ismaster.arbiters;
        self.tags = ismaster.tags;
        self.set_name = ismaster.set_name;
        self.election_id = ismaster.election_id;
        self.primary = ismaster.primary;

        let hosts_empty = self.hosts.is_empty();
        let set_name_empty = self.set_name.is_empty();
        let msg_empty = ismaster.msg.is_empty();

        self.stype = if msg_empty && set_name_empty && hosts_empty {
            ServerType::Standalone
        } else if !msg_empty {
            ServerType::Mongos
        } else if ismaster.is_master && !set_name_empty {
            ServerType::RSPrimary
        } else if ismaster.is_secondary && !set_name_empty {
            ServerType::RSSecondary
        } else if ismaster.arbiter_only && !set_name_empty {
            ServerType::RSArbiter
        } else if !set_name_empty {
            ServerType::RSOther
        } else if ismaster.is_replica_set {
            ServerType::RSGhost
        } else {
            ServerType::Unknown
        }
    }

    pub fn set_err(&mut self, err: Error) {
        self.err = Arc::new(Some(err));
        self.stype = ServerType::Unknown;
    }
}

impl Monitor {
    pub fn new(host: Host, description: Arc<RwLock<ServerDescription>>) -> Result<Monitor> {
        Ok(Monitor {
            socket: try!(TcpStream::connect((&host.host_name[..], host.port))),
            host: host,
            description: description,
        })
    }

    pub fn reconnect(&mut self) -> Result<()> {
        let ref host_name = self.host.host_name;
        let port = self.host.port;
        self.socket = try!(TcpStream::connect((&host_name[..], port)));
        Ok(())
    }
}

impl Server {
    pub fn new(req_id: Arc<AtomicIsize>, host: Host) -> Server {
        let description = Arc::new(RwLock::new(ServerDescription::new()));

        // Create new monitor thread
        let host_clone = host.clone();
        let desc_clone = description.clone();

        thread::spawn(move|| {
            let mut monitor = Monitor::new(host_clone, desc_clone).ok().expect("Failed to connect monitor to server.");

            // Call ismaster on socket at low level to avoid using client resources
            let options = FindOptions::new().with_limit(1);
            let flags = OpQueryFlags::with_find_options(&options);
            let mut filter = bson::Document::new();
            filter.insert("isMaster".to_owned(), Bson::I32(1));

            loop {
                // break on some condition somehow

                let result = Cursor::query_with_socket(
                    &mut monitor.socket, None, req_id.fetch_add(1, Ordering::SeqCst) as i32,
                    "local.$cmd".to_owned(), options.batch_size, flags, options.skip as i32,
                    options.limit, filter.clone(), options.projection.clone(), false);

                match result {
                    Ok(mut cursor) => {
                        match cursor.next() {
                            Some(Ok(doc)) => {
                                // Parse ismaster result and update server description.
                                let ismaster_result = IsMasterResult::new(doc);
                                {
                                    let mut server_description = monitor.description.write().unwrap();
                                    match ismaster_result {
                                        Ok(ismaster) => server_description.update(ismaster),
                                        Err(err) => server_description.set_err(err),
                                    }
                                }
                            },
                            Some(Err(err)) => panic!(err),
                            None => panic!("ismaster returned no response."),
                        }
                        thread::sleep_ms(DEFAULT_HEARTBEAT_FREQUENCY_MS);
                    },
                    Err(err) => {
                        {
                            let mut server_description = monitor.description.write().unwrap();
                            server_description.set_err(err);
                        }
                        thread::sleep_ms(DEFAULT_RETRY_FREQUENCY_MS);
                    }
                }
            }
        });

        Server {
            host: host.clone(),
            pool: ConnectionPool::new(host),
            description: description.clone(),
        }
    }

    pub fn acquire_stream(&self) -> Result<PooledStream> {
        self.pool.acquire_stream()
    }
}
<commit_msg>Add Monitor run cmd and retry<commit_after>use Error::{self, ArgumentError, OperationError};
use Result;

use bson::{self, Bson, oid};
use chrono::{DateTime, UTC};

use coll::options::FindOptions;
use connstring::{self, ConnectionString, Host};
use cursor::Cursor;
use pool::{ConnectionPool, PooledStream};
use wire_protocol::flags::OpQueryFlags;

use std::collections::{BTreeMap, HashMap};
use std::net::TcpStream;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicIsize, Ordering};
use std::thread;

const DEFAULT_HEARTBEAT_FREQUENCY_MS: u32 = 10000;
const DEFAULT_RETRY_FREQUENCY_MS: u32 = 1000;
const DEFAULT_MAX_BSON_OBJECT_SIZE: i64 = 16 * 1024 * 1024;
const DEFAULT_MAX_MESSAGE_SIZE_BYTES: i64 = 48000000;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TopologyType {
    Single,
    ReplicaSetNoPrimary,
    ReplicaSetWithPrimary,
    Sharded,
    Unknown,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ServerType {
    Standalone,
    Mongos,
    RSPrimary,
    RSSecondary,
    RSArbiter,
    RSOther,
    RSGhost,
    Unknown,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct IsMasterResult {
    is_master: bool,
    max_bson_object_size: i64,
    max_message_size_bytes: i64,
    local_time: DateTime<UTC>,
    min_wire_version: i64,
    max_wire_version: i64,

    // Shards
    msg: String,

    // RS
    is_replica_set: bool,
    is_secondary: bool,
    me: Option<Host>,
    hosts: Vec<Host>,
    passives: Vec<Host>,
    arbiters: Vec<Host>,
    arbiter_only: bool,
    tags: BTreeMap<String, String>,
    set_name: String,
    election_id: Option<oid::ObjectId>,
    primary: Option<Host>,
    hidden: bool,
}

#[derive(Clone)]
pub struct Topology {
    config: ConnectionString,
    description: TopologyDescription,
    servers: Arc<HashMap<Host, RwLock<Server>>>,
    compatible: bool,
    compat_error: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TopologyDescription {
    ttype: TopologyType,
    set_name: String,
    heartbeat_frequency_ms: u32,
    max_election_id: Option<oid::ObjectId>,
}

#[derive(Clone)]
pub struct Server {
    pub host: Host,
    pool: ConnectionPool,
    description: Arc<RwLock<ServerDescription>>,
}

#[derive(Clone, Debug)]
pub struct ServerDescription {
    stype: ServerType,
    err: Arc<Option<Error>>,
    round_trip_time: Option<i64>,
    min_wire_version: i64,
    max_wire_version: i64,
    me: Option<Host>,
    hosts: Vec<Host>,
    passives: Vec<Host>,
    arbiters: Vec<Host>,
    tags: BTreeMap<String, String>,
    set_name: String,
    election_id: Option<oid::ObjectId>,
    primary: Option<Host>,
}

struct Monitor {
    host: Host,
    description: Arc<RwLock<ServerDescription>>,
    socket: TcpStream,
    req_id: Arc<AtomicIsize>,
    running: bool,
}

impl IsMasterResult {
    pub fn new(doc: bson::Document) -> Result<IsMasterResult> {
        let is_master = match doc.get("ismaster") {
            Some(&Bson::Boolean(ref b)) => *b,
            _ => return Err(ArgumentError("result does not contain 'ismaster'.".to_owned())),
        };

        let local_time = match doc.get("localTime") {
            Some(&Bson::UtcDatetime(ref datetime)) => datetime.clone(),
            _ => return Err(ArgumentError("result does not contain 'localTime'.".to_owned())),
        };

        let min_version = match doc.get("minWireVersion") {
            Some(&Bson::I64(ref v)) => *v,
            _ => return Err(ArgumentError("result does not contain 'minWireVersion'.".to_owned())),
        };

        let max_version = match doc.get("maxWireVersion") {
            Some(&Bson::I64(ref v)) => *v,
            _ => return Err(ArgumentError("result does not contain 'maxWireVersion'.".to_owned())),
        };

        let mut result = IsMasterResult {
            is_master: is_master,
            max_bson_object_size: DEFAULT_MAX_BSON_OBJECT_SIZE,
            max_message_size_bytes: DEFAULT_MAX_MESSAGE_SIZE_BYTES,
            local_time: local_time,
            min_wire_version: min_version,
            max_wire_version: max_version,
            msg: String::new(),
            is_secondary: false,
            is_replica_set: false,
            me: None,
            hosts: Vec::new(),
            passives: Vec::new(),
            arbiters: Vec::new(),
            arbiter_only: false,
            tags: BTreeMap::new(),
            set_name: String::new(),
            election_id: None,
            primary: None,
            hidden: false,
        };

        if let Some(&Bson::String(ref s)) = doc.get("msg") {
            result.msg = s.to_owned();
        }

        if let Some(&Bson::Boolean(ref b)) = doc.get("secondary") {
            result.is_secondary = *b;
        }

        if let Some(&Bson::Boolean(ref b)) = doc.get("isreplicaset") {
            result.is_replica_set = *b;
        }

        if let Some(&Bson::String(ref s)) = doc.get("setName") {
            result.set_name = s.to_owned();
        }

        if let Some(&Bson::String(ref s)) = doc.get("me") {
            result.me = Some(try!(connstring::parse_host(s)));
        }

        if let Some(&Bson::Array(ref arr)) = doc.get("hosts") {
            result.hosts = arr.iter().filter_map(|bson| match bson {
                &Bson::String(ref s) => connstring::parse_host(s).ok(),
                _ => None,
            }).collect();
        }

        if let Some(&Bson::Array(ref arr)) = doc.get("passives") {
            result.passives = arr.iter().filter_map(|bson| match bson {
                &Bson::String(ref s) => connstring::parse_host(s).ok(),
                _ => None,
            }).collect();
        }

        if let Some(&Bson::Array(ref arr)) = doc.get("arbiters") {
            result.passives = arr.iter().filter_map(|bson| match bson {
                &Bson::String(ref s) => connstring::parse_host(s).ok(),
                _ => None,
            }).collect();
        }

        if let Some(&Bson::String(ref s)) = doc.get("primary") {
            result.primary = Some(try!(connstring::parse_host(s)));
        }

        if let Some(&Bson::Boolean(ref b)) = doc.get("arbiterOnly") {
            result.arbiter_only = *b;
        }

        if let Some(&Bson::Boolean(ref h)) = doc.get("hidden") {
            result.hidden = *h;
        }

        if let Some(&Bson::Document(ref doc)) = doc.get("tags") {
            for (k, v) in doc.into_iter() {
                if let &Bson::String(ref tag) = v {
                    result.tags.insert(k.to_owned(), tag.to_owned());
                }
            }
        }

        if let Some(&Bson::ObjectId(ref id)) = doc.get("electionId") {
            result.election_id = Some(id.clone());
        }

        Ok(result)
    }
}

impl TopologyDescription {
    pub fn new() -> TopologyDescription {
        TopologyDescription {
            ttype: TopologyType::Unknown,
            set_name: String::new(),
            heartbeat_frequency_ms: DEFAULT_HEARTBEAT_FREQUENCY_MS,
            max_election_id: None,
        }
    }
}

impl Topology {
    pub fn new(req_id: Arc<AtomicIsize>, config: ConnectionString,
               description: Option<TopologyDescription>) -> Result<Topology> {

        let options = description.unwrap_or(TopologyDescription::new());

        if config.hosts.len() > 1 && options.ttype == TopologyType::Single {
            return Err(ArgumentError(
                "TopologyType::Single cannot be used with multiple seeds.".to_owned()));
        }

        if !options.set_name.is_empty() && options.ttype != TopologyType::ReplicaSetNoPrimary {
            return Err(ArgumentError(
                "TopologyType must be ReplicaSetNoPrimary if set_name is provided.".to_owned()));
        }

        // TODO: Determine driver's wire compatibility, and check overlap with
        // all servers in topology.

        let mut servers = HashMap::new();
        for host in config.hosts.iter() {
            let server = Server::new(req_id.clone(), host.clone());
            servers.insert(host.clone(), RwLock::new(server));
        }

        Ok(Topology {
            config: config,
            description: options,
            servers: Arc::new(servers),
            compatible: true,
            compat_error: String::new(),
        })
    }

    pub fn acquire_stream(&self) -> Result<PooledStream> {
        for (_, server) in self.servers.iter() {
            let read_server = try!(server.read());
            return read_server.acquire_stream()
        }
        Err(OperationError("No servers found in configuration.".to_owned()))
    }
}

impl ServerDescription {
    pub fn new() -> ServerDescription {
        ServerDescription {
            stype: ServerType::Unknown,
            err: Arc::new(None),
            round_trip_time: None,
            min_wire_version: 0,
            max_wire_version: 0,
            me: None,
            hosts: Vec::new(),
            passives: Vec::new(),
            arbiters: Vec::new(),
            tags: BTreeMap::new(),
            set_name: String::new(),
            election_id: None,
            primary: None,
        }
    }

    pub fn update(&mut self, ismaster: IsMasterResult) {
        self.min_wire_version = ismaster.min_wire_version;
        self.max_wire_version = ismaster.max_wire_version;
        self.me = ismaster.me;
        self.hosts = ismaster.hosts;
        self.passives = ismaster.passives;
        self.arbiters = ismaster.arbiters;
        self.tags = ismaster.tags;
        self.set_name = ismaster.set_name;
        self.election_id = ismaster.election_id;
        self.primary = ismaster.primary;

        let hosts_empty = self.hosts.is_empty();
        let set_name_empty = self.set_name.is_empty();
        let msg_empty = ismaster.msg.is_empty();

        self.stype = if msg_empty && set_name_empty && hosts_empty {
            ServerType::Standalone
        } else if !msg_empty {
            ServerType::Mongos
        } else if ismaster.is_master && !set_name_empty {
            ServerType::RSPrimary
        } else if ismaster.is_secondary && !set_name_empty {
            ServerType::RSSecondary
        } else if ismaster.arbiter_only && !set_name_empty {
            ServerType::RSArbiter
        } else if !set_name_empty {
            ServerType::RSOther
        } else if ismaster.is_replica_set {
            ServerType::RSGhost
        } else {
            ServerType::Unknown
        }
    }

    pub fn set_err(&mut self, err: Error) {
        self.err = Arc::new(Some(err));
        self.stype = ServerType::Unknown;
    }
}

impl Monitor {
    pub fn new(host: Host, description: Arc<RwLock<ServerDescription>>,
               req_id: Arc<AtomicIsize>) -> Result<Monitor> {

        Ok(Monitor {
            socket: try!(TcpStream::connect((&host.host_name[..], host.port))),
            req_id: req_id,
            host: host,
            description: description,
            running: false,
        })
    }

    pub fn reconnect(&mut self) -> Result<()> {
        let ref host_name = self.host.host_name;
        let port = self.host.port;
        self.socket = try!(TcpStream::connect((&host_name[..], port)));
        Ok(())
    }

    pub fn update_server_description(& self, doc: bson::Document) {
        // Parse ismaster result and update server description.
        let ismaster_result = IsMasterResult::new(doc);
        let mut server_description = self.description.write().unwrap();
        match ismaster_result {
            Ok(ismaster) => server_description.update(ismaster),
            Err(err) => server_description.set_err(err),
        }
    }

    pub fn run(&mut self) {
        if self.running {
            return;
        }

        self.running = true;
        // Call ismaster on socket at low level to avoid using client resources
        let options = FindOptions::new().with_limit(1);
        let flags = OpQueryFlags::with_find_options(&options);
        let mut filter = bson::Document::new();
        filter.insert("isMaster".to_owned(), Bson::I32(1));

        loop {
            if !self.running {
                break;
            }

            let result = Cursor::query_with_socket(
                &mut self.socket, None, self.req_id.fetch_add(1, Ordering::SeqCst) as i32,
                "local.$cmd".to_owned(), options.batch_size, flags, options.skip as i32,
                options.limit, filter.clone(), options.projection.clone(), false);

            match result {
                Ok(mut cursor) => {
                    match cursor.next() {
                        Some(Ok(doc)) => self.update_server_description(doc),
                        Some(Err(err)) => panic!(err),
                        None => panic!("ismaster returned no response."),
                    }
                    thread::sleep_ms(DEFAULT_HEARTBEAT_FREQUENCY_MS);
                },
                Err(err) => {
                    // clear connection pool for server here

                    if self.description.read().unwrap().stype == ServerType::Unknown {
                        let mut server_description = self.description.write().unwrap();
                        server_description.set_err(err);
                    } else {
                        // Retry once
                        let result = Cursor::query_with_socket(
                            &mut self.socket, None, self.req_id.fetch_add(1, Ordering::SeqCst) as i32,
                            "local.$cmd".to_owned(), options.batch_size, flags, options.skip as i32,
                            options.limit, filter.clone(), options.projection.clone(), false);

                        match result {
                            Ok(mut cursor) => {
                                match cursor.next() {
                                    Some(Ok(doc)) => self.update_server_description(doc),
                                    Some(Err(err)) => panic!(err),
                                    None => panic!("ismaster returned no response."),
                                }
                            },
                            Err(err) => {
                                let mut server_description = self.description.write().unwrap();
                                server_description.set_err(err);
                                server_description.stype = ServerType::Unknown;
                            }
                        }
                    }
                    thread::sleep_ms(DEFAULT_RETRY_FREQUENCY_MS);
                }
            }
        }
    }
}

impl Server {
    pub fn new(req_id: Arc<AtomicIsize>, host: Host) -> Server {
        let description = Arc::new(RwLock::new(ServerDescription::new()));

        // Create new monitor thread
        let host_clone = host.clone();
        let desc_clone = description.clone();

        let mut monitor = Monitor::new(host_clone, desc_clone, req_id).ok().expect("Failed to connect monitor to server.");

        thread::spawn(move || {
            monitor.run();
        });

        Server {
            host: host.clone(),
            pool: ConnectionPool::new(host),
            description: description.clone(),
        }
    }

    pub fn acquire_stream(&self) -> Result<PooledStream> {
        self.pool.acquire_stream()
    }
}
<|endoftext|>


<commit_before>// This file is part of Grust, GObject introspection bindings for Rust
//
// Copyright (C) 2014  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

use types::gpointer;
use wrap;
use wrap::Wrapper;

use std::mem;
use std::ops::{Deref, DerefMut};

pub type RefFunc = unsafe extern "C" fn(gpointer) -> gpointer;
pub type UnrefFunc = unsafe extern "C" fn(gpointer);
pub type RefcountFuncs = (*const RefFunc, *const UnrefFunc);

pub trait Refcount {
    fn refcount_funcs(&self) -> &'static RefcountFuncs;
}

pub struct Ref<T: Refcount> {
    plumbing: RefImpl
}

pub struct SyncRef<T: Refcount + Send + Sync> {
    plumbing: RefImpl
}

unsafe impl<T> Send for SyncRef<T> where T: Refcount + Send + Sync { }
unsafe impl<T> Sync for SyncRef<T> where T: Refcount + Send + Sync { }

impl<T> Ref<T> where T: Refcount {

    pub fn new(source: &mut T) -> Ref<T> {
        new_ref(source)
    }
}

impl<T> Ref<T> where T: Refcount + Wrapper {

    pub unsafe fn from_raw(ptr: *mut <T as Wrapper>::Raw) -> Ref<T> {
        Ref::new(wrap::from_raw_mut(ptr, &ptr))
    }
}

impl<T> SyncRef<T> where T: Refcount + Send + Sync {

    pub fn new(source: &mut T) -> SyncRef<T> {
        new_sync_ref(source)
    }
}

impl<T> SyncRef<T> where T: Refcount + Wrapper + Send + Sync {

    pub unsafe fn from_raw(ptr: *mut <T as Wrapper>::Raw) -> SyncRef<T> {
        SyncRef::new(wrap::from_raw_mut(ptr, &ptr))
    }
}

struct RefImpl {
    ptr: gpointer
}

impl RefImpl {

    unsafe fn impl_drop(&mut self, (_, dec_ref): RefcountFuncs) {
        (*dec_ref)(self.ptr);
    }

    unsafe fn impl_clone_from(&mut self, source: &RefImpl,
                              (inc_ref, dec_ref): RefcountFuncs)
    {
        if self.ptr == source.ptr {
            return;
        }
        (*inc_ref)(source.ptr);
        (*dec_ref)(self.ptr);
        self.ptr = source.ptr;
    }
}

unsafe fn new_ref_impl(p: gpointer, (inc_ref, _): RefcountFuncs) -> RefImpl {
    (*inc_ref)(p);
    RefImpl { ptr: p }
}

fn new_ref<T>(source: &T) -> Ref<T>
    where T: Refcount
{
    let p = source as *const T as gpointer;
    let funcs = source.refcount_funcs();
    Ref {
        plumbing: unsafe { new_ref_impl(p, *funcs) }
    }
}

fn new_sync_ref<T>(source: &T) -> SyncRef<T>
    where T: Refcount + Send + Sync
{
    let p = source as *const T as gpointer;
    let funcs = source.refcount_funcs();
    SyncRef {
        plumbing: unsafe { new_ref_impl(p, *funcs) }
    }
}

#[unsafe_destructor]
impl<T> Drop for Ref<T> where T: Refcount {
    fn drop(&mut self) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_drop(*funcs);
        }
    }
}

#[unsafe_destructor]
impl<T> Drop for SyncRef<T> where T: Refcount + Send + Sync {
    fn drop(&mut self) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_drop(*funcs);
        }
    }
}

impl<T: Refcount> Clone for Ref<T> {

    fn clone(&self) -> Ref<T> {
        new_ref(self.deref())
    }

    fn clone_from(&mut self, source: &Ref<T>) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_clone_from(&source.plumbing, *funcs);
        }
    }
}

impl<T: Refcount + Send + Sync> Clone for SyncRef<T> {

    fn clone(&self) -> SyncRef<T> {
        new_sync_ref(self.deref())
    }

    fn clone_from(&mut self, source: &SyncRef<T>) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_clone_from(&source.plumbing, *funcs);
        }
    }
}

impl<T> Deref for Ref<T> where T: Refcount {

    type Target = T;

    fn deref(&self) -> &T {
        let p = self.plumbing.ptr as *const T;
        unsafe { mem::copy_lifetime(self, &*p) }
    }
}

impl<T> Deref for SyncRef<T> where T: Refcount + Send + Sync {

    type Target = T;

    fn deref(&self) -> &T {
        let p = self.plumbing.ptr as *const T;
        unsafe { mem::copy_lifetime(self, &*p) }
    }
}

impl<T> DerefMut for Ref<T> where T: Refcount {

    fn deref_mut(&mut self) -> &mut T {
        let p = self.plumbing.ptr as *mut T;
        unsafe { mem::copy_mut_lifetime(self, &mut *p) }
    }
}

impl<T> DerefMut for SyncRef<T> where T: Refcount + Send + Sync {

    fn deref_mut(&mut self) -> &mut T {
        let p = self.plumbing.ptr as *mut T;
        unsafe { mem::copy_mut_lifetime(self, &mut *p) }
    }
}
<commit_msg>refcount: Don't leak a reference in from_raw constructors<commit_after>// This file is part of Grust, GObject introspection bindings for Rust
//
// Copyright (C) 2014  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

use types::gpointer;
use wrap::Wrapper;

use std::mem;
use std::ops::{Deref, DerefMut};

pub type RefFunc = unsafe extern "C" fn(gpointer) -> gpointer;
pub type UnrefFunc = unsafe extern "C" fn(gpointer);
pub type RefcountFuncs = (*const RefFunc, *const UnrefFunc);

pub trait Refcount {
    fn refcount_funcs(&self) -> &'static RefcountFuncs;
}

pub struct Ref<T: Refcount> {
    plumbing: RefImpl
}

pub struct SyncRef<T: Refcount + Send + Sync> {
    plumbing: RefImpl
}

unsafe impl<T> Send for SyncRef<T> where T: Refcount + Send + Sync { }
unsafe impl<T> Sync for SyncRef<T> where T: Refcount + Send + Sync { }

impl<T> Ref<T> where T: Refcount {

    pub fn new(source: &mut T) -> Ref<T> {
        new_ref(source)
    }
}

impl<T> Ref<T> where T: Refcount + Wrapper {

    pub unsafe fn from_raw(ptr: *mut <T as Wrapper>::Raw) -> Ref<T> {
        Ref {
            plumbing: RefImpl { ptr: ptr as gpointer }
        }
    }
}

impl<T> SyncRef<T> where T: Refcount + Send + Sync {

    pub fn new(source: &mut T) -> SyncRef<T> {
        new_sync_ref(source)
    }
}

impl<T> SyncRef<T> where T: Refcount + Wrapper + Send + Sync {

    pub unsafe fn from_raw(ptr: *mut <T as Wrapper>::Raw) -> SyncRef<T> {
        SyncRef {
            plumbing: RefImpl { ptr: ptr as gpointer }
        }
    }
}

struct RefImpl {
    ptr: gpointer
}

impl RefImpl {

    unsafe fn impl_drop(&mut self, (_, dec_ref): RefcountFuncs) {
        (*dec_ref)(self.ptr);
    }

    unsafe fn impl_clone_from(&mut self, source: &RefImpl,
                              (inc_ref, dec_ref): RefcountFuncs)
    {
        if self.ptr == source.ptr {
            return;
        }
        (*inc_ref)(source.ptr);
        (*dec_ref)(self.ptr);
        self.ptr = source.ptr;
    }
}

unsafe fn new_ref_impl(p: gpointer, (inc_ref, _): RefcountFuncs) -> RefImpl {
    (*inc_ref)(p);
    RefImpl { ptr: p }
}

fn new_ref<T>(source: &T) -> Ref<T>
    where T: Refcount
{
    let p = source as *const T as gpointer;
    let funcs = source.refcount_funcs();
    Ref {
        plumbing: unsafe { new_ref_impl(p, *funcs) }
    }
}

fn new_sync_ref<T>(source: &T) -> SyncRef<T>
    where T: Refcount + Send + Sync
{
    let p = source as *const T as gpointer;
    let funcs = source.refcount_funcs();
    SyncRef {
        plumbing: unsafe { new_ref_impl(p, *funcs) }
    }
}

#[unsafe_destructor]
impl<T> Drop for Ref<T> where T: Refcount {
    fn drop(&mut self) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_drop(*funcs);
        }
    }
}

#[unsafe_destructor]
impl<T> Drop for SyncRef<T> where T: Refcount + Send + Sync {
    fn drop(&mut self) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_drop(*funcs);
        }
    }
}

impl<T: Refcount> Clone for Ref<T> {

    fn clone(&self) -> Ref<T> {
        new_ref(self.deref())
    }

    fn clone_from(&mut self, source: &Ref<T>) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_clone_from(&source.plumbing, *funcs);
        }
    }
}

impl<T: Refcount + Send + Sync> Clone for SyncRef<T> {

    fn clone(&self) -> SyncRef<T> {
        new_sync_ref(self.deref())
    }

    fn clone_from(&mut self, source: &SyncRef<T>) {
        let funcs = self.refcount_funcs();
        unsafe {
            self.plumbing.impl_clone_from(&source.plumbing, *funcs);
        }
    }
}

impl<T> Deref for Ref<T> where T: Refcount {

    type Target = T;

    fn deref(&self) -> &T {
        let p = self.plumbing.ptr as *const T;
        unsafe { mem::copy_lifetime(self, &*p) }
    }
}

impl<T> Deref for SyncRef<T> where T: Refcount + Send + Sync {

    type Target = T;

    fn deref(&self) -> &T {
        let p = self.plumbing.ptr as *const T;
        unsafe { mem::copy_lifetime(self, &*p) }
    }
}

impl<T> DerefMut for Ref<T> where T: Refcount {

    fn deref_mut(&mut self) -> &mut T {
        let p = self.plumbing.ptr as *mut T;
        unsafe { mem::copy_mut_lifetime(self, &mut *p) }
    }
}

impl<T> DerefMut for SyncRef<T> where T: Refcount + Send + Sync {

    fn deref_mut(&mut self) -> &mut T {
        let p = self.plumbing.ptr as *mut T;
        unsafe { mem::copy_mut_lifetime(self, &mut *p) }
    }
}
<|endoftext|>