<commit_before>extern crate rocksdb;
extern crate rustc_serialize;

use std::collections::HashMap;

use self::rustc_serialize::json::{JsonEvent, Parser, StackElement};
// Needed for a trait in order to `put()` into a `rocksdb::WriteBatch`
use self::rocksdb::Writable;

use key_builder::{KeyBuilder, SegmentType};
use records_capnp::payload;
use stems::Stems;
use query::{Error};

// Good example of using rustc_serialize: https://github.com/ajroetker/beautician/blob/master/src/lib.rs
// Callback based JSON streaming parser: https://github.com/gyscos/json-streamer.rs
// Another parser pased on rustc_serializ: https://github.com/isagalaev/ijson-rust/blob/master/src/test.rs#L11

#[derive(Debug, PartialEq)]
struct WordInfo {
    //offset in the text field where the stemmed text starts
    stemmed_offset: u64,

    // the suffix of the stemmed text. When applied over stemmed, the original
    // text is returned.
    suffix_text: String,

    // the start of the suffixText
    suffix_offset: u64,
}

type ArrayOffsets = Vec<u64>;
type ArrayOffsetsToWordInfo = HashMap<ArrayOffsets, Vec<WordInfo>>;
type WordPathInfoMap = HashMap<String, ArrayOffsetsToWordInfo>;

enum ObjectKeyTypes {
    /// _id field
    Id,
    /// Normal key
    Key(String),
    /// Reserved key starting with underscore
    Ignore,
    /// No key found
    NoKey,
}

#[derive(Debug)]
pub struct Shredder {
    keybuilder: KeyBuilder,
    map: WordPathInfoMap,
    path_array_offsets: ArrayOffsets,
    // Top-level fields prefixed with an underscore are ignored
    ignore_children: u64,
    doc_id: String,
}

impl Shredder {
    pub fn new() -> Shredder {
        Shredder{
            keybuilder: KeyBuilder::new(),
            map: WordPathInfoMap::new(),
            path_array_offsets: Vec::new(),
            ignore_children: 0,
            doc_id: String::new(),
        }
    }
    fn add_entries(&mut self, text: &String, docseq: u64) {
        let stems = Stems::new(text.as_str());
        for stem in stems {
            self.keybuilder.push_word(&stem.stemmed);
            self.keybuilder.push_doc_seq(docseq);
            let map_path_array_offsets = self.map.entry(self.keybuilder.key())
                                                        .or_insert(ArrayOffsetsToWordInfo::new());
            let map_word_infos = map_path_array_offsets.entry(self.path_array_offsets.clone())
                .or_insert(Vec::new());
            map_word_infos.push(WordInfo{
                stemmed_offset: stem.stemmed_offset as u64,
                suffix_text: stem.suffix.to_string(),
                suffix_offset: stem.suffix_offset as u64,
            });
            self.keybuilder.pop_doc_seq();
            self.keybuilder.pop_word();
        }
        println!("add_entries: map: {:?}", self.map);
    }

    fn inc_top_array_offset(&mut self) {
        // we encounter a new element. if we are a child element of an array
        // increment the offset. If we aren't (we are the root value or a map
        // value) we don't increment
        if let Some(SegmentType::Array) = self.keybuilder.last_pushed_segment_type() {
            if let Some(last) = self.path_array_offsets.last_mut() {
                *last += 1;
            }
        }
    }

    // Extract key if it exists and indicates if it's a special type of key
    fn extract_key(&mut self, stack_element: Option<StackElement>) -> ObjectKeyTypes {
        if self.keybuilder.last_pushed_segment_type().unwrap() == SegmentType::ObjectKey {
            match stack_element {
                Some(StackElement::Key(key)) => {
                    if self.keybuilder.segments.len() == 1 && key.starts_with("_") {
                        if key == "_id" {
                            ObjectKeyTypes::Id
                        } else {
                            ObjectKeyTypes::Ignore
                        }
                    } else {
                        ObjectKeyTypes::Key(key.to_string())
                    }
                },
                _ => ObjectKeyTypes::NoKey,
            }
        } else {
            ObjectKeyTypes::NoKey
        }
    }

    // If we are inside an object we need to push the key to the key builder
    // Don't push them if they are reserved fields (starting with underscore)
    fn maybe_push_key(&mut self, stack_element: Option<StackElement>) -> Result<(), Error> {
        if self.keybuilder.last_pushed_segment_type().unwrap() == SegmentType::ObjectKey
                && self.keybuilder.segments.len() == 1 {
            if let Some(StackElement::Key(key)) = stack_element {
                if key.starts_with("_") {
                    if key == "_id" {
                        return Err(Error::Shred(
                            "Expected string for `_id` field, got another type".to_string()));
                    } else {
                        self.ignore_children = 1;
                    }
                } else {
                    // Pop the dummy object that makes ObjectEnd happy
                    // or the previous object key
                    self.keybuilder.pop_object_key();
                    self.keybuilder.push_object_key(key.to_string());
                }
            }
        }
        Ok(())
    }

   pub fn shred(&mut self, json: &str, docseq: u64) -> Result<String, Error> {
        let mut parser = Parser::new(json.chars());
        let mut token = parser.next();

        loop {
            // Get the next token, so that in case of an `ObjectStart` the key is already
            // on the stack.
            match token.take() {
                Some(JsonEvent::ObjectStart) => {
                    if self.ignore_children > 0 {
                        self.ignore_children += 1;
                    }
                    else {
                        // Just push something to make `ObjectEnd` happy
                        self.keybuilder.push_object_key("".to_string());
                    }
                },
                Some(JsonEvent::ObjectEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {//if !self.keybuilder.segments.is_empty() {
                        self.keybuilder.pop_object_key();
                    }
                },
                Some(JsonEvent::ArrayStart) => {
                    if self.ignore_children > 0 {
                        self.ignore_children += 1;
                    } else {
                        self.keybuilder.push_array();
                        self.path_array_offsets.push(0);
                    }
                },
                Some(JsonEvent::ArrayEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {
                        self.path_array_offsets.pop();
                        self.keybuilder.pop_array();
                    }
                },
                Some(JsonEvent::StringValue(value)) => {
                    // No children to ignore
                    if self.ignore_children == 0 {
                        println!("stringvalue: {:?}", value);
                        match self.extract_key(parser.stack().top()) {
                            ObjectKeyTypes::Id => self.doc_id = value,
                            ObjectKeyTypes::Key(key) => {
                                // Pop the dummy object that makes ObjectEnd happy
                                // or the previous object key
                                self.keybuilder.pop_object_key();
                                self.keybuilder.push_object_key(key.to_string());

                                self.add_entries(&value, docseq);
                                self.inc_top_array_offset();
                            },
                            ObjectKeyTypes::NoKey => {
                                self.add_entries(&value, docseq);
                                self.inc_top_array_offset();
                            },
                            ObjectKeyTypes::Ignore => {
                                self.ignore_children = 1;
                            },
                        }
                    }
                },
                not_implemented => {
                    panic!("Not yet implemented other JSON types! {:?}", not_implemented);
                }
            };

            token = parser.next();
            if token == None {
                break;
            }
        }
        println!("keybuilder: {}", self.keybuilder.key());
        println!("shredder: keys:");
        for key in self.map.keys() {
            println!("  {}", key);
        }
        Ok(self.doc_id.clone())
    }

    pub fn add_to_batch(&self, batch: &rocksdb::WriteBatch) -> Result<(), Error> {
        for (key_path, word_path_infos) in &self.map {
            let mut message = ::capnp::message::Builder::new_default();
            {
                let capn_payload = message.init_root::<payload::Builder>();
                let mut capn_arrayoffsets_to_wordinfo = capn_payload.init_arrayoffsets_to_wordinfos(
                    word_path_infos.len() as u32);
                for (infos_pos, (arrayoffsets, wordinfos)) in word_path_infos.iter().enumerate() {

                    let mut capn_a2w = capn_arrayoffsets_to_wordinfo.borrow().get(infos_pos as u32);
                    {
                        let mut capn_arrayoffsets =
                            capn_a2w.borrow().init_arrayoffsets(arrayoffsets.len() as u32);
                        for (pos, arrayoffset) in arrayoffsets.iter().enumerate() {
                            capn_arrayoffsets.set(pos as u32, arrayoffset.clone());
                        }
                    }
                    {
                        let mut capn_wordinfos = capn_a2w.init_wordinfos(wordinfos.len() as u32);
                        for (pos, wordinfo) in wordinfos.iter().enumerate() {
                            let mut capn_wordinfo = capn_wordinfos.borrow().get(pos as u32);
                            capn_wordinfo.set_stemmed_offset(wordinfo.stemmed_offset);
                            capn_wordinfo.set_suffix_text(&wordinfo.suffix_text);
                            capn_wordinfo.set_suffix_offset(wordinfo.suffix_offset);
                        }
                    }
                }
            }
            let mut bytes = Vec::new();
            ::capnp::serialize_packed::write_message(&mut bytes, &message).unwrap();
            try!(batch.put(&key_path.clone().into_bytes(), &bytes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{ArrayOffsetsToWordInfo, WordInfo, WordPathInfoMap};

    #[test]
    fn test_shred_nested() {
        let mut shredder = super::Shredder::new();
        //let json = r#"{"hello": {"my": "world!"}, "anumber": 2}"#;
        //let json = r#"{"A":[{"B":"B2VMX two three","C":"C2"},{"B": "b1","C":"C2"}]}"#;
        //let json = r#"{"A":[[[{"B": "string within deeply nested array should be stemmed"}]]]}"#;
        //let json = r#"[{"A": 1, "B": 2, "C": 3}]"#;
        //let json = r#"{"foo": {"bar": 1}}"#;
        let json = r#"{"some": ["array", "data", ["also", "nested"]]}"#;
        let docseq = 123;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.some$!array#123", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.some$!data#123", vec![
                (vec![1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!also#123", vec![
                (vec![2, 0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!nest#123", vec![
                (vec![2, 1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "ed".to_string(), suffix_offset: 4 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    #[test]
    fn test_shred_objects() {
        let mut shredder = super::Shredder::new();
        let json = r#"{"A":[{"B":"B2VMX two three","C":"..C2"},{"B": "b1","C":"..C2"}]}"#;
        let docseq = 1234;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.A$.B!b1#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }])]),
            ("W.A$.B!b2vmx#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.A$.B!c2#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 },
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.A$.B!three#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 10, suffix_text: "".to_string(), suffix_offset: 15 }])]),
            ("W.A$.B!two#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 6, suffix_text: "".to_string(), suffix_offset: 9 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    fn compare_shredded(result_map: &WordPathInfoMap,
                        expected: &Vec<(&str, Vec<(Vec<u64>, Vec<WordInfo>)>)>) {
        // HashMap have an arbitraty order of the elements
        let mut result: Vec<(&String, &ArrayOffsetsToWordInfo)> = result_map.into_iter().collect();
        result.sort_by(|a, b| Ord::cmp(&a.0, &b.0));
        for (ii, &(key, values)) in result.iter().enumerate() {
            assert_eq!(key, expected[ii].0);
            let mut wordinfos: Vec<(&Vec<u64>, &Vec<WordInfo>)> = values.iter().collect();
            wordinfos.sort_by_key(|item| item.0);
            for (jj, wordinfo) in wordinfos.iter().enumerate() {
                assert_eq!(wordinfo.0, &expected[ii].1[jj].0);
                assert_eq!(wordinfo.1, &expected[ii].1[jj].1);
            }
        }
    }

    #[test]
    fn test_shred_empty_object() {
        let mut shredder = super::Shredder::new();
        let json = r#"{}"#;
        let docseq = 123;
        shredder.shred(json, docseq).unwrap();
        // TODO vmx 2016-11-08: Finish this test
    }
}
<commit_msg>Fix tests<commit_after>extern crate rocksdb;
extern crate rustc_serialize;

use std::collections::HashMap;

use self::rustc_serialize::json::{JsonEvent, Parser, StackElement};
// Needed for a trait in order to `put()` into a `rocksdb::WriteBatch`
use self::rocksdb::Writable;

use key_builder::{KeyBuilder, SegmentType};
use records_capnp::payload;
use stems::Stems;
use query::{Error};

// Good example of using rustc_serialize: https://github.com/ajroetker/beautician/blob/master/src/lib.rs
// Callback based JSON streaming parser: https://github.com/gyscos/json-streamer.rs
// Another parser pased on rustc_serializ: https://github.com/isagalaev/ijson-rust/blob/master/src/test.rs#L11

#[derive(Debug, PartialEq)]
struct WordInfo {
    //offset in the text field where the stemmed text starts
    stemmed_offset: u64,

    // the suffix of the stemmed text. When applied over stemmed, the original
    // text is returned.
    suffix_text: String,

    // the start of the suffixText
    suffix_offset: u64,
}

type ArrayOffsets = Vec<u64>;
type ArrayOffsetsToWordInfo = HashMap<ArrayOffsets, Vec<WordInfo>>;
type WordPathInfoMap = HashMap<String, ArrayOffsetsToWordInfo>;

enum ObjectKeyTypes {
    /// _id field
    Id,
    /// Normal key
    Key(String),
    /// Reserved key starting with underscore
    Ignore,
    /// No key found
    NoKey,
}

#[derive(Debug)]
pub struct Shredder {
    keybuilder: KeyBuilder,
    map: WordPathInfoMap,
    path_array_offsets: ArrayOffsets,
    // Top-level fields prefixed with an underscore are ignored
    ignore_children: u64,
    doc_id: String,
}

impl Shredder {
    pub fn new() -> Shredder {
        Shredder{
            keybuilder: KeyBuilder::new(),
            map: WordPathInfoMap::new(),
            path_array_offsets: Vec::new(),
            ignore_children: 0,
            doc_id: String::new(),
        }
    }
    fn add_entries(&mut self, text: &String, docseq: u64) {
        let stems = Stems::new(text.as_str());
        for stem in stems {
            self.keybuilder.push_word(&stem.stemmed);
            self.keybuilder.push_doc_seq(docseq);
            let map_path_array_offsets = self.map.entry(self.keybuilder.key())
                                                        .or_insert(ArrayOffsetsToWordInfo::new());
            let map_word_infos = map_path_array_offsets.entry(self.path_array_offsets.clone())
                .or_insert(Vec::new());
            map_word_infos.push(WordInfo{
                stemmed_offset: stem.stemmed_offset as u64,
                suffix_text: stem.suffix.to_string(),
                suffix_offset: stem.suffix_offset as u64,
            });
            self.keybuilder.pop_doc_seq();
            self.keybuilder.pop_word();
        }
        println!("add_entries: map: {:?}", self.map);
    }

    fn inc_top_array_offset(&mut self) {
        // we encounter a new element. if we are a child element of an array
        // increment the offset. If we aren't (we are the root value or a map
        // value) we don't increment
        if let Some(SegmentType::Array) = self.keybuilder.last_pushed_segment_type() {
            if let Some(last) = self.path_array_offsets.last_mut() {
                *last += 1;
            }
        }
    }

    // Extract key if it exists and indicates if it's a special type of key
    fn extract_key(&mut self, stack_element: Option<StackElement>) -> ObjectKeyTypes {
        if self.keybuilder.last_pushed_segment_type().unwrap() == SegmentType::ObjectKey {
            match stack_element {
                Some(StackElement::Key(key)) => {
                    if self.keybuilder.segments.len() == 1 && key.starts_with("_") {
                        if key == "_id" {
                            ObjectKeyTypes::Id
                        } else {
                            ObjectKeyTypes::Ignore
                        }
                    } else {
                        ObjectKeyTypes::Key(key.to_string())
                    }
                },
                _ => ObjectKeyTypes::NoKey,
            }
        } else {
            ObjectKeyTypes::NoKey
        }
    }

    // If we are inside an object we need to push the key to the key builder
    // Don't push them if they are reserved fields (starting with underscore)
    fn maybe_push_key(&mut self, stack_element: Option<StackElement>) -> Result<(), Error> {
        if self.keybuilder.last_pushed_segment_type().unwrap() == SegmentType::ObjectKey
                && self.keybuilder.segments.len() == 1 {
            if let Some(StackElement::Key(key)) = stack_element {
                if key.starts_with("_") {
                    if key == "_id" {
                        return Err(Error::Shred(
                            "Expected string for `_id` field, got another type".to_string()));
                    } else {
                        self.ignore_children = 1;
                    }
                } else {
                    // Pop the dummy object that makes ObjectEnd happy
                    // or the previous object key
                    self.keybuilder.pop_object_key();
                    self.keybuilder.push_object_key(key.to_string());
                }
            }
        }
        Ok(())
    }

   pub fn shred(&mut self, json: &str, docseq: u64) -> Result<String, Error> {
        let mut parser = Parser::new(json.chars());
        let mut token = parser.next();

        loop {
            // Get the next token, so that in case of an `ObjectStart` the key is already
            // on the stack.
            match token.take() {
                Some(JsonEvent::ObjectStart) => {
                    if self.ignore_children > 0 {
                        self.ignore_children += 1;
                    }
                    else {
                        // Just push something to make `ObjectEnd` happy
                        self.keybuilder.push_object_key("".to_string());
                    }
                },
                Some(JsonEvent::ObjectEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {//if !self.keybuilder.segments.is_empty() {
                        self.keybuilder.pop_object_key();
                    }
                },
                Some(JsonEvent::ArrayStart) => {
                    try!(self.maybe_push_key(parser.stack().top()));
                    if self.ignore_children > 0 {
                        self.ignore_children += 1;
                    } else {
                        self.keybuilder.push_array();
                        self.path_array_offsets.push(0);
                    }
                },
                Some(JsonEvent::ArrayEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {
                        self.path_array_offsets.pop();
                        self.keybuilder.pop_array();
                    }
                },
                Some(JsonEvent::StringValue(value)) => {
                    // No children to ignore
                    if self.ignore_children == 0 {
                        println!("stringvalue: {:?}", value);
                        match self.extract_key(parser.stack().top()) {
                            ObjectKeyTypes::Id => self.doc_id = value,
                            ObjectKeyTypes::Key(key) => {
                                // Pop the dummy object that makes ObjectEnd happy
                                // or the previous object key
                                self.keybuilder.pop_object_key();
                                self.keybuilder.push_object_key(key.to_string());

                                self.add_entries(&value, docseq);
                                self.inc_top_array_offset();
                            },
                            ObjectKeyTypes::NoKey => {
                                self.add_entries(&value, docseq);
                                self.inc_top_array_offset();
                            },
                            ObjectKeyTypes::Ignore => {
                                self.ignore_children = 1;
                            },
                        }
                    }
                },
                not_implemented => {
                    panic!("Not yet implemented other JSON types! {:?}", not_implemented);
                }
            };

            token = parser.next();
            if token == None {
                break;
            }
        }
        println!("keybuilder: {}", self.keybuilder.key());
        println!("shredder: keys:");
        for key in self.map.keys() {
            println!("  {}", key);
        }
        Ok(self.doc_id.clone())
    }

    pub fn add_to_batch(&self, batch: &rocksdb::WriteBatch) -> Result<(), Error> {
        for (key_path, word_path_infos) in &self.map {
            let mut message = ::capnp::message::Builder::new_default();
            {
                let capn_payload = message.init_root::<payload::Builder>();
                let mut capn_arrayoffsets_to_wordinfo = capn_payload.init_arrayoffsets_to_wordinfos(
                    word_path_infos.len() as u32);
                for (infos_pos, (arrayoffsets, wordinfos)) in word_path_infos.iter().enumerate() {

                    let mut capn_a2w = capn_arrayoffsets_to_wordinfo.borrow().get(infos_pos as u32);
                    {
                        let mut capn_arrayoffsets =
                            capn_a2w.borrow().init_arrayoffsets(arrayoffsets.len() as u32);
                        for (pos, arrayoffset) in arrayoffsets.iter().enumerate() {
                            capn_arrayoffsets.set(pos as u32, arrayoffset.clone());
                        }
                    }
                    {
                        let mut capn_wordinfos = capn_a2w.init_wordinfos(wordinfos.len() as u32);
                        for (pos, wordinfo) in wordinfos.iter().enumerate() {
                            let mut capn_wordinfo = capn_wordinfos.borrow().get(pos as u32);
                            capn_wordinfo.set_stemmed_offset(wordinfo.stemmed_offset);
                            capn_wordinfo.set_suffix_text(&wordinfo.suffix_text);
                            capn_wordinfo.set_suffix_offset(wordinfo.suffix_offset);
                        }
                    }
                }
            }
            let mut bytes = Vec::new();
            ::capnp::serialize_packed::write_message(&mut bytes, &message).unwrap();
            try!(batch.put(&key_path.clone().into_bytes(), &bytes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{ArrayOffsetsToWordInfo, WordInfo, WordPathInfoMap};

    #[test]
    fn test_shred_nested() {
        let mut shredder = super::Shredder::new();
        //let json = r#"{"hello": {"my": "world!"}, "anumber": 2}"#;
        //let json = r#"{"A":[{"B":"B2VMX two three","C":"C2"},{"B": "b1","C":"C2"}]}"#;
        //let json = r#"{"A":[[[{"B": "string within deeply nested array should be stemmed"}]]]}"#;
        //let json = r#"[{"A": 1, "B": 2, "C": 3}]"#;
        //let json = r#"{"foo": {"bar": 1}}"#;
        let json = r#"{"some": ["array", "data", ["also", "nested"]]}"#;
        let docseq = 123;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.some$!array#123", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.some$!data#123", vec![
                (vec![1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!also#123", vec![
                (vec![2, 0], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ("W.some$$!nest#123", vec![
                (vec![2, 1], vec![WordInfo {
                    stemmed_offset: 0, suffix_text: "ed".to_string(), suffix_offset: 4 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    #[test]
    fn test_shred_objects() {
        let mut shredder = super::Shredder::new();
        let json = r#"{"A":[{"B":"B2VMX two three","C":"..C2"},{"B": "b1","C":"..C2"}]}"#;
        let docseq = 1234;
        shredder.shred(json, docseq).unwrap();
        let expected = vec![
            ("W.A$.B!b1#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }])]),
            ("W.A$.B!b2vmx#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }])]),
            ("W.A$.B!three#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 10, suffix_text: "".to_string(), suffix_offset: 15 }])]),
            ("W.A$.B!two#1234", vec![
                (vec![0], vec![WordInfo {
                    stemmed_offset: 6, suffix_text: "".to_string(), suffix_offset: 9 }])]),
            ("W.A$.C!c2#1234", vec![
                (vec![0], vec![
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 },
                    WordInfo {
                        stemmed_offset: 2, suffix_text: "".to_string(), suffix_offset: 4 }])]),
            ];
        compare_shredded(&shredder.map, &expected);
    }

    fn compare_shredded(result_map: &WordPathInfoMap,
                        expected: &Vec<(&str, Vec<(Vec<u64>, Vec<WordInfo>)>)>) {
        // HashMap have an arbitraty order of the elements
        let mut result: Vec<(&String, &ArrayOffsetsToWordInfo)> = result_map.into_iter().collect();
        result.sort_by(|a, b| Ord::cmp(&a.0, &b.0));
        for (ii, &(key, values)) in result.iter().enumerate() {
            assert_eq!(key, expected[ii].0);
            let mut wordinfos: Vec<(&Vec<u64>, &Vec<WordInfo>)> = values.iter().collect();
            wordinfos.sort_by_key(|item| item.0);
            for (jj, wordinfo) in wordinfos.iter().enumerate() {
                assert_eq!(wordinfo.0, &expected[ii].1[jj].0);
                assert_eq!(wordinfo.1, &expected[ii].1[jj].1);
            }
        }
    }

    #[test]
    fn test_shred_empty_object() {
        let mut shredder = super::Shredder::new();
        let json = r#"{}"#;
        let docseq = 123;
        shredder.shred(json, docseq).unwrap();
        // TODO vmx 2016-11-08: Finish this test
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;

use typed_arena::Arena;

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

struct Thing<'a> {
    node_arena: &'a Arena<Node<'a>>,
    leaves: Vec<Leaf>,
    lists: Vec<List<'a>>,
}

#[derive(Debug)]
struct Node<'a> {
    weight: usize,
    count: usize,
    tail: Option<&'a Node<'a>>,
}

#[derive(Debug)]
struct Leaf {
    weight: usize,
    count: usize,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug, Clone)]
struct List<'a> {
    lookahead0: &'a Node<'a>,
    lookahead1: &'a Node<'a>,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    let node0 = node_arena.alloc(Node {
        weight: leaves[0].weight,
        count: 1,
        tail: None,
    });

    let node1 = node_arena.alloc(Node {
        weight: leaves[1].weight,
        count: 2,
        tail: None,
    });

    let lists: Vec<List> = vec![
        List {
            lookahead0: node0,
            lookahead1: node1,
        };
        max_bits
    ];

    let mut thing = Thing {
        node_arena: &node_arena,
        leaves,
        lists,
    };

    // In the last list, 2 * num_symbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * num_symbols - 4;
    for _ in 0..num_boundary_pm_runs - 1 {
        thing.boundary_pm(max_bits - 1);
    }

    thing.boundary_pm_final(max_bits - 1);

    // let mut a = lists.pop().unwrap().lookahead1.leaf_counts.into_iter().rev().peekable();
    //
    // let mut bitlength_value = 1;
    // while let Some(leaf_count) = a.next() {
    //     let next_count = *a.peek().unwrap_or(&0);
    //     for leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

impl<'a> Thing<'a> {
    fn boundary_pm(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        let last_count = self.lists[index].lookahead1.count; // Count of last chain of list.

        if index == 0 && last_count >= num_symbols {
            return
        }

        self.lists[index].lookahead0 = self.lists[index].lookahead1;

        if index == 0 {
            // New leaf node in list 0.
            let new_chain = self.node_arena.alloc(Node {
               weight: self.leaves[last_count].weight,
               count: last_count + 1,
               tail: self.lists[index].lookahead0.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            let weight_sum = {
                let previous_list = &self.lists[index - 1];
                previous_list.lookahead0.weight + previous_list.lookahead1.weight
            };

            if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
                // New leaf inserted in list, so count is incremented.
                let new_chain = self.node_arena.alloc(Node {
                   weight: self.leaves[last_count].weight,
                   count: last_count + 1,
                   tail: self.lists[index].lookahead0.tail,
                });
                self.lists[index].lookahead1 = new_chain;
            } else {
                let new_chain = self.node_arena.alloc(Node {
                   weight: weight_sum,
                   count: last_count,
                   tail: Some(self.lists[index - 1].lookahead1),
                });
                self.lists[index].lookahead1 = new_chain;

                // Two lookahead chains of previous list used up, create new ones.
                self.boundary_pm(index - 1);
                self.boundary_pm(index - 1);
            }
        }
    }

    fn boundary_pm_final(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        // Count of last chain of list.
        let last_count = self.lists[index].lookahead1.count;

        let weight_sum = {
            let previous_list = &self.lists[index - 1];
            previous_list.lookahead0.weight + previous_list.lookahead1.weight
        };

        if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
            let new_chain = self.node_arena.alloc(Node {
               weight: 0,
               count: last_count + 1,
               tail: self.lists[index].lookahead1.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            // ???
            self.lists[index].lookahead1.tail = Some(self.lists[index - 1].lookahead1);
        }
    }
}

// fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let mut current_list = &mut lists[current_list_index];
//
//     // The next leaf goes next; counting itself makes the leaf_count increase by one.
//     current_list.lookahead1.weight = leaves[current_list.next_leaf_index].weight;
//     current_list.lookahead1.leaf_counts.clear();
//     current_list.lookahead1.leaf_counts.extend_from_slice(&current_list.lookahead0.leaf_counts);
//     let last_index = current_list.lookahead1.leaf_counts.len() - 1;
//     current_list.lookahead1.leaf_counts[last_index] += 1;
//     current_list.next_leaf_index += 1;
// }

// fn next_tree(weight_sum: usize, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     {
//         let (head, tail) = lists.split_at_mut(current_list_index);
//         let prev_list = head.last_mut().unwrap();
//         let current_list = tail.first_mut().unwrap();
//
//         let previous_list_leaf_counts = &prev_list.lookahead1.leaf_counts;
//
//         // Make a tree from the lookaheads from the previous list; that goes next.
//         // This is not a leaf node, so the leaf count stays the same.
//         current_list.lookahead1.weight = weight_sum;
//         current_list.lookahead1.leaf_counts.clear();
//
//         current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
//         current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.leaf_counts.last().unwrap());
//     }
//
//     // The previous list needs two new lookahead nodes.
//     boundary_pm(lists, leaves, current_list_index - 1);
//     boundary_pm(lists, leaves, current_list_index - 1);
// }

// fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
//     let last_index = lists.len() - 1;
//     boundary_pm(lists, leaves, last_index);
// }

// fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let next_leaf_index = lists[current_list_index].next_leaf_index;
//     let num_symbols = leaves.len();
//
//     if current_list_index == 0 && next_leaf_index == num_symbols {
//         // We've added all the leaves to the lowest list, so we're done here
//         return;
//     }
//
//     mem::swap(&mut lists[current_list_index].lookahead0, &mut lists[current_list_index].lookahead1);
//
//     if current_list_index == 0 {
//         lowest_list(lists, leaves);
//     } else {
//         // We're at a list other than the lowest list.
//         let weight_sum = {
//             let previous_list = &lists[current_list_index - 1];
//             previous_list.lookahead0.weight + previous_list.lookahead1.weight
//         };
//
//         if next_leaf_index < num_symbols && weight_sum > leaves[next_leaf_index].weight {
//             next_leaf(lists, leaves, current_list_index);
//         } else {
//             next_tree(weight_sum, lists, leaves, current_list_index);
//         }
//     }
// }

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<commit_msg>Use Cell to be able to change the tail<commit_after>use std::cell::Cell;
use std::cmp::Ordering;

use typed_arena::Arena;

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

struct Thing<'a> {
    node_arena: &'a Arena<Node<'a>>,
    leaves: Vec<Leaf>,
    lists: Vec<List<'a>>,
}

#[derive(Debug)]
struct Node<'a> {
    weight: usize,
    count: usize,
    tail: Cell<Option<&'a Node<'a>>>,
}

#[derive(Debug)]
struct Leaf {
    weight: usize,
    count: usize,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug, Clone)]
struct List<'a> {
    lookahead0: &'a Node<'a>,
    lookahead1: &'a Node<'a>,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    let node0 = node_arena.alloc(Node {
        weight: leaves[0].weight,
        count: 1,
        tail: Cell::new(None),
    });

    let node1 = node_arena.alloc(Node {
        weight: leaves[1].weight,
        count: 2,
        tail: Cell::new(None),
    });

    let lists: Vec<List> = vec![
        List {
            lookahead0: node0,
            lookahead1: node1,
        };
        max_bits
    ];

    let mut thing = Thing {
        node_arena: &node_arena,
        leaves,
        lists,
    };

    // In the last list, 2 * num_symbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * num_symbols - 4;
    for _ in 0..num_boundary_pm_runs - 1 {
        thing.boundary_pm(max_bits - 1);
    }

    thing.boundary_pm_final(max_bits - 1);

    // let mut a = lists.pop().unwrap().lookahead1.leaf_counts.into_iter().rev().peekable();
    //
    // let mut bitlength_value = 1;
    // while let Some(leaf_count) = a.next() {
    //     let next_count = *a.peek().unwrap_or(&0);
    //     for leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

impl<'a> Thing<'a> {
    fn boundary_pm(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        let last_count = self.lists[index].lookahead1.count; // Count of last chain of list.

        if index == 0 && last_count >= num_symbols {
            return
        }

        self.lists[index].lookahead0 = self.lists[index].lookahead1;

        if index == 0 {
            // New leaf node in list 0.
            let new_chain = self.node_arena.alloc(Node {
               weight: self.leaves[last_count].weight,
               count: last_count + 1,
               tail: self.lists[index].lookahead0.tail.clone(),
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            let weight_sum = {
                let previous_list = &self.lists[index - 1];
                previous_list.lookahead0.weight + previous_list.lookahead1.weight
            };

            if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
                // New leaf inserted in list, so count is incremented.
                let new_chain = self.node_arena.alloc(Node {
                   weight: self.leaves[last_count].weight,
                   count: last_count + 1,
                   tail: self.lists[index].lookahead0.tail.clone(),
                });
                self.lists[index].lookahead1 = new_chain;
            } else {
                let new_chain = self.node_arena.alloc(Node {
                   weight: weight_sum,
                   count: last_count,
                   tail: Cell::new(Some(self.lists[index - 1].lookahead1)),
                });
                self.lists[index].lookahead1 = new_chain;

                // Two lookahead chains of previous list used up, create new ones.
                self.boundary_pm(index - 1);
                self.boundary_pm(index - 1);
            }
        }
    }

    fn boundary_pm_final(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        // Count of last chain of list.
        let last_count = self.lists[index].lookahead1.count;

        let weight_sum = {
            let previous_list = &self.lists[index - 1];
            previous_list.lookahead0.weight + previous_list.lookahead1.weight
        };

        if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
            let new_chain = self.node_arena.alloc(Node {
               weight: 0,
               count: last_count + 1,
               tail: self.lists[index].lookahead1.tail.clone(),
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            self.lists[index].lookahead1.tail.set(Some(self.lists[index - 1].lookahead1));
        }
    }
}

// fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let mut current_list = &mut lists[current_list_index];
//
//     // The next leaf goes next; counting itself makes the leaf_count increase by one.
//     current_list.lookahead1.weight = leaves[current_list.next_leaf_index].weight;
//     current_list.lookahead1.leaf_counts.clear();
//     current_list.lookahead1.leaf_counts.extend_from_slice(&current_list.lookahead0.leaf_counts);
//     let last_index = current_list.lookahead1.leaf_counts.len() - 1;
//     current_list.lookahead1.leaf_counts[last_index] += 1;
//     current_list.next_leaf_index += 1;
// }

// fn next_tree(weight_sum: usize, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     {
//         let (head, tail) = lists.split_at_mut(current_list_index);
//         let prev_list = head.last_mut().unwrap();
//         let current_list = tail.first_mut().unwrap();
//
//         let previous_list_leaf_counts = &prev_list.lookahead1.leaf_counts;
//
//         // Make a tree from the lookaheads from the previous list; that goes next.
//         // This is not a leaf node, so the leaf count stays the same.
//         current_list.lookahead1.weight = weight_sum;
//         current_list.lookahead1.leaf_counts.clear();
//
//         current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
//         current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.leaf_counts.last().unwrap());
//     }
//
//     // The previous list needs two new lookahead nodes.
//     boundary_pm(lists, leaves, current_list_index - 1);
//     boundary_pm(lists, leaves, current_list_index - 1);
// }

// fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
//     let last_index = lists.len() - 1;
//     boundary_pm(lists, leaves, last_index);
// }

// fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let next_leaf_index = lists[current_list_index].next_leaf_index;
//     let num_symbols = leaves.len();
//
//     if current_list_index == 0 && next_leaf_index == num_symbols {
//         // We've added all the leaves to the lowest list, so we're done here
//         return;
//     }
//
//     mem::swap(&mut lists[current_list_index].lookahead0, &mut lists[current_list_index].lookahead1);
//
//     if current_list_index == 0 {
//         lowest_list(lists, leaves);
//     } else {
//         // We're at a list other than the lowest list.
//         let weight_sum = {
//             let previous_list = &lists[current_list_index - 1];
//             previous_list.lookahead0.weight + previous_list.lookahead1.weight
//         };
//
//         if next_leaf_index < num_symbols && weight_sum > leaves[next_leaf_index].weight {
//             next_leaf(lists, leaves, current_list_index);
//         } else {
//             next_tree(weight_sum, lists, leaves, current_list_index);
//         }
//     }
// }

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<|endoftext|>


<commit_before>//! Store XML element attributes and their values.

use std::slice;
use std::str;

use markup5ever::{namespace_url, LocalName, Namespace, Prefix, QualName};
use string_cache::DefaultAtom;

use crate::util::{opt_utf8_cstr, utf8_cstr};

/// Type used to store attribute values.
///
/// Attribute values are often repeated in an SVG file, so we intern them using the
/// string_cache crate.
pub type AttributeValue = DefaultAtom;

/// Iterable wrapper for libxml2's representation of attribute/value.
///
/// See the [`new_from_xml2_attributes`] function for information.
///
/// [`new_from_xml2_attributes`]: #method.new_from_xml2_attributes
#[derive(Clone)]
pub struct Attributes(Vec<(QualName, AttributeValue)>);

/// Iterator from `Attributes.iter`.
pub struct AttributesIter<'a>(slice::Iter<'a, (QualName, AttributeValue)>);

impl Attributes {
    /// Creates an iterable `Attributes` from the C array of borrowed C strings.
    ///
    /// With libxml2's SAX parser, the caller's startElementNsSAX2Func
    /// callback gets passed a `xmlChar **` for attributes, which
    /// comes in groups of (localname/prefix/URI/value_start/value_end).
    /// In those, localname/prefix/URI are NUL-terminated strings;
    /// value_start and value_end point to the start-inclusive and
    /// end-exclusive bytes in the attribute's value.
    ///
    /// # Safety
    ///
    /// This function is unsafe because the caller must guarantee the following:
    ///
    /// * `attrs` is a valid pointer, with (n_attributes * 5) elements.
    ///
    /// * All strings are valid UTF-8.
    pub unsafe fn new_from_xml2_attributes(
        n_attributes: usize,
        attrs: *const *const libc::c_char,
    ) -> Attributes {
        let mut array = Vec::new();

        if n_attributes > 0 && !attrs.is_null() {
            for attr in slice::from_raw_parts(attrs, n_attributes * 5).chunks_exact(5) {
                let localname = attr[0];
                let prefix = attr[1];
                let uri = attr[2];
                let value_start = attr[3];
                let value_end = attr[4];

                assert!(!localname.is_null());

                let localname = utf8_cstr(localname);

                let prefix = opt_utf8_cstr(prefix);
                let uri = opt_utf8_cstr(uri);
                let qual_name = QualName::new(
                    prefix.map(Prefix::from),
                    uri.map(Namespace::from)
                        .unwrap_or_else(|| namespace_url!("")),
                    LocalName::from(localname),
                );

                if !value_start.is_null() && !value_end.is_null() {
                    assert!(value_end >= value_start);

                    // FIXME: ptr::offset_from() is nightly-only.
                    // We'll do the computation of the length by hand.
                    let start = value_start as usize;
                    let end = value_end as usize;
                    let len = end - start;

                    let value_slice = slice::from_raw_parts(value_start as *const u8, len);
                    let value_str = str::from_utf8_unchecked(value_slice);
                    let value_atom = DefaultAtom::from(value_str);

                    array.push((qual_name, value_atom));
                }
            }
        }

        Attributes(array)
    }

    /// Returns the number of attributes.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Creates an iterator that yields `(QualName, &'a str)` tuples.
    pub fn iter(&self) -> AttributesIter<'_> {
        AttributesIter(self.0.iter())
    }
}

impl<'a> Iterator for AttributesIter<'a> {
    type Item = (QualName, &'a str);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(a, v)| (a.clone(), v.as_ref()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use markup5ever::{expanded_name, local_name, namespace_url, ns};
    use std::ffi::CString;
    use std::ptr;

    #[test]
    fn empty_attributes() {
        let map = unsafe { Attributes::new_from_xml2_attributes(0, ptr::null()) };
        assert_eq!(map.len(), 0);
    }

    #[test]
    fn attributes_with_namespaces() {
        let attrs = [
            (
                CString::new("href").unwrap(),
                Some(CString::new("xlink").unwrap()),
                Some(CString::new("http://www.w3.org/1999/xlink").unwrap()),
                CString::new("1").unwrap(),
            ),
            (
                CString::new("ry").unwrap(),
                None,
                None,
                CString::new("2").unwrap(),
            ),
            (
                CString::new("d").unwrap(),
                None,
                None,
                CString::new("").unwrap(),
            ),
        ];

        let mut v: Vec<*const libc::c_char> = Vec::new();

        for (localname, prefix, uri, val) in &attrs {
            v.push(localname.as_ptr() as *const libc::c_char);
            v.push(
                prefix
                    .as_ref()
                    .map(|p: &CString| p.as_ptr())
                    .unwrap_or_else(|| ptr::null()) as *const libc::c_char,
            );
            v.push(
                uri.as_ref()
                    .map(|p: &CString| p.as_ptr())
                    .unwrap_or_else(|| ptr::null()) as *const libc::c_char,
            );

            let val_start = val.as_ptr() as *const libc::c_char;
            let val_end = unsafe { val_start.offset(val.as_bytes().len() as isize) };
            v.push(val_start); // value_start
            v.push(val_end); // value_end
        }

        let attrs = unsafe { Attributes::new_from_xml2_attributes(3, v.as_ptr()) };

        let mut had_href: bool = false;
        let mut had_ry: bool = false;
        let mut had_d: bool = false;

        for (a, v) in attrs.iter() {
            match a.expanded() {
                expanded_name!(xlink "href") => {
                    assert!(v == "1");
                    had_href = true;
                }

                expanded_name!("", "ry") => {
                    assert!(v == "2");
                    had_ry = true;
                }

                expanded_name!("", "d") => {
                    assert!(v == "");
                    had_d = true;
                }

                _ => unreachable!(),
            }
        }

        assert!(had_href);
        assert!(had_ry);
        assert!(had_d);
    }
}
<commit_msg>Trivial casts in tests<commit_after>//! Store XML element attributes and their values.

use std::slice;
use std::str;

use markup5ever::{namespace_url, LocalName, Namespace, Prefix, QualName};
use string_cache::DefaultAtom;

use crate::util::{opt_utf8_cstr, utf8_cstr};

/// Type used to store attribute values.
///
/// Attribute values are often repeated in an SVG file, so we intern them using the
/// string_cache crate.
pub type AttributeValue = DefaultAtom;

/// Iterable wrapper for libxml2's representation of attribute/value.
///
/// See the [`new_from_xml2_attributes`] function for information.
///
/// [`new_from_xml2_attributes`]: #method.new_from_xml2_attributes
#[derive(Clone)]
pub struct Attributes(Vec<(QualName, AttributeValue)>);

/// Iterator from `Attributes.iter`.
pub struct AttributesIter<'a>(slice::Iter<'a, (QualName, AttributeValue)>);

impl Attributes {
    /// Creates an iterable `Attributes` from the C array of borrowed C strings.
    ///
    /// With libxml2's SAX parser, the caller's startElementNsSAX2Func
    /// callback gets passed a `xmlChar **` for attributes, which
    /// comes in groups of (localname/prefix/URI/value_start/value_end).
    /// In those, localname/prefix/URI are NUL-terminated strings;
    /// value_start and value_end point to the start-inclusive and
    /// end-exclusive bytes in the attribute's value.
    ///
    /// # Safety
    ///
    /// This function is unsafe because the caller must guarantee the following:
    ///
    /// * `attrs` is a valid pointer, with (n_attributes * 5) elements.
    ///
    /// * All strings are valid UTF-8.
    pub unsafe fn new_from_xml2_attributes(
        n_attributes: usize,
        attrs: *const *const libc::c_char,
    ) -> Attributes {
        let mut array = Vec::new();

        if n_attributes > 0 && !attrs.is_null() {
            for attr in slice::from_raw_parts(attrs, n_attributes * 5).chunks_exact(5) {
                let localname = attr[0];
                let prefix = attr[1];
                let uri = attr[2];
                let value_start = attr[3];
                let value_end = attr[4];

                assert!(!localname.is_null());

                let localname = utf8_cstr(localname);

                let prefix = opt_utf8_cstr(prefix);
                let uri = opt_utf8_cstr(uri);
                let qual_name = QualName::new(
                    prefix.map(Prefix::from),
                    uri.map(Namespace::from)
                        .unwrap_or_else(|| namespace_url!("")),
                    LocalName::from(localname),
                );

                if !value_start.is_null() && !value_end.is_null() {
                    assert!(value_end >= value_start);

                    // FIXME: ptr::offset_from() is nightly-only.
                    // We'll do the computation of the length by hand.
                    let start = value_start as usize;
                    let end = value_end as usize;
                    let len = end - start;

                    let value_slice = slice::from_raw_parts(value_start as *const u8, len);
                    let value_str = str::from_utf8_unchecked(value_slice);
                    let value_atom = DefaultAtom::from(value_str);

                    array.push((qual_name, value_atom));
                }
            }
        }

        Attributes(array)
    }

    /// Returns the number of attributes.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Creates an iterator that yields `(QualName, &'a str)` tuples.
    pub fn iter(&self) -> AttributesIter<'_> {
        AttributesIter(self.0.iter())
    }
}

impl<'a> Iterator for AttributesIter<'a> {
    type Item = (QualName, &'a str);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(a, v)| (a.clone(), v.as_ref()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use markup5ever::{expanded_name, local_name, namespace_url, ns};
    use std::ffi::CString;
    use std::ptr;

    #[test]
    fn empty_attributes() {
        let map = unsafe { Attributes::new_from_xml2_attributes(0, ptr::null()) };
        assert_eq!(map.len(), 0);
    }

    #[test]
    fn attributes_with_namespaces() {
        let attrs = [
            (
                CString::new("href").unwrap(),
                Some(CString::new("xlink").unwrap()),
                Some(CString::new("http://www.w3.org/1999/xlink").unwrap()),
                CString::new("1").unwrap(),
            ),
            (
                CString::new("ry").unwrap(),
                None,
                None,
                CString::new("2").unwrap(),
            ),
            (
                CString::new("d").unwrap(),
                None,
                None,
                CString::new("").unwrap(),
            ),
        ];

        let mut v: Vec<*const libc::c_char> = Vec::new();

        for (localname, prefix, uri, val) in &attrs {
            v.push(localname.as_ptr());
            v.push(
                prefix
                    .as_ref()
                    .map(|p: &CString| p.as_ptr())
                    .unwrap_or_else(|| ptr::null()),
            );
            v.push(
                uri.as_ref()
                    .map(|p: &CString| p.as_ptr())
                    .unwrap_or_else(|| ptr::null()),
            );

            let val_start = val.as_ptr();
            let val_end = unsafe { val_start.offset(val.as_bytes().len() as isize) };
            v.push(val_start); // value_start
            v.push(val_end); // value_end
        }

        let attrs = unsafe { Attributes::new_from_xml2_attributes(3, v.as_ptr()) };

        let mut had_href: bool = false;
        let mut had_ry: bool = false;
        let mut had_d: bool = false;

        for (a, v) in attrs.iter() {
            match a.expanded() {
                expanded_name!(xlink "href") => {
                    assert!(v == "1");
                    had_href = true;
                }

                expanded_name!("", "ry") => {
                    assert!(v == "2");
                    had_ry = true;
                }

                expanded_name!("", "d") => {
                    assert!(v == "");
                    had_d = true;
                }

                _ => unreachable!(),
            }
        }

        assert!(had_href);
        assert!(had_ry);
        assert!(had_d);
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2016 Élisabeth HENRY.
//
// This file is part of Crowbook.
//
// Crowbook 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.
//
// Caribon 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 ba copy of the GNU Lesser General Public License
// along with Crowbook.  If not, see <http://www.gnu.org/licenses/>.

//! Provides utility functions for escaping text to HTML or LaTeX

use std::borrow::Cow;

const NB_CHAR:char = ' '; // non breaking space
const NB_CHAR_NARROW:char = '\u{202F}'; // narrow non breaking space
const NB_CHAR_EM:char = '\u{2002}'; // demi em space

/// Escape non breaking spaces for HTML, so they are visible.
#[doc(hidden)]
pub fn escape_nb_spaces<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    let input = input.into();
    if input.contains(|c| match c {
        NB_CHAR | NB_CHAR_NARROW | NB_CHAR_EM => true,
        _ => false
    }) {
        let mut output = String::with_capacity(input.len());
        for c in input.chars() {
            match c {
                NB_CHAR_NARROW
                    | NB_CHAR_EM
                    | NB_CHAR
                    => output.push_str(&format!("{}",
                                                match c {
                                                    NB_CHAR => "#ffff66",
                                                    NB_CHAR_NARROW => "#9999ff",
                                                    NB_CHAR_EM => "#ff9999",
                                                    _ => unreachable!()
                                                },
                                                c)),
                _ => output.push(c),
            }
        }
        Cow::Owned(output)
    } else {
        input.into()
    }
}

/// Escape characters `<`, `>`, and `&`
///
/// # Examples
///
/// ```
/// use crowbook::escape::escape_html;
/// let s = escape_html("<foo> & <bar>");
/// assert_eq!(&s, "<foo> & <bar>");
/// ```
pub fn escape_html<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    let input = input.into();
    if input.contains(|c| match c {
        '<'|'>'|'&' => true,
        _ => false
    }) {
        let mut output = String::with_capacity(input.len());
        for c in input.chars() {
            match c {
                '<' => output.push_str("<"),
                '>' => output.push_str(">"),
                '&' => output.push_str("&"),
                _ => output.push(c),
            }
        }

        Cow::Owned(output)
    } else {
        input.into()
    }
}

/// Escape characters for LaTeX
///
/// # Examples
///
/// ```
/// use crowbook::escape::escape_tex;
/// let s = escape_tex("command --foo # calls command with option foo");
/// assert_eq!(&s, r"command -{}-foo \# calls command with option foo");
/// ```
pub fn escape_tex<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    let input = input.into();
    if input.contains(|c| match c {
        '&'|'%'|'$'|'#'|'_'|'{'|'}'|'~'|'^'|'\\'|'-' => true,
        _ => false
    }) {
        let mut output = String::with_capacity(input.len());
        let mut chars:Vec<char> = input.chars().collect();
        chars.push(' '); // add a dummy char for call to .windows()
        // for &[c, next] in chars.windows(2) { // still experimental, uncomment when stable
        for win in chars.windows(2) {
            let c = win[0];
            let next = win[1];
            match c {
                '-' => {
                    if next == '-' {
                        output.push_str(r"-{}"); // if next char is also a -, to avoid tex ligatures
                    } else {
                        output.push(c);
                    }
                },
                '&' => output.push_str(r"\&"),
                '%' => output.push_str(r"\%"),
                '$' => output.push_str(r"\$"),
                '#' => output.push_str(r"\#"),
                '_' => output.push_str(r"\_"),
                '{' => output.push_str(r"\{"),
                '}' => output.push_str(r"\}"),
                '~' => output.push_str(r"\textasciitilde{}" ),
                '^' => output.push_str(r"\textasciicircum{}"),
                '\\' => output.push_str(r"\textbackslash{}"),
                _  => output.push(c)
            }
        }
        Cow::Owned(output)
    } else {
        input
    }
}
<commit_msg>Try to optimize a little escape_html<commit_after>// Copyright (C) 2016 Élisabeth HENRY.
//
// This file is part of Crowbook.
//
// Crowbook 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.
//
// Caribon 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 ba copy of the GNU Lesser General Public License
// along with Crowbook.  If not, see <http://www.gnu.org/licenses/>.

//! Provides utility functions for escaping text to HTML or LaTeX

use std::borrow::Cow;

const NB_CHAR:char = ' '; // non breaking space
const NB_CHAR_NARROW:char = '\u{202F}'; // narrow non breaking space
const NB_CHAR_EM:char = '\u{2002}'; // demi em space

/// Escape non breaking spaces for HTML, so they are visible.
#[doc(hidden)]
pub fn escape_nb_spaces<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    let input = input.into();
    if input.contains(|c| match c {
        NB_CHAR | NB_CHAR_NARROW | NB_CHAR_EM => true,
        _ => false
    }) {
        let mut output = String::with_capacity(input.len());
        for c in input.chars() {
            match c {
                NB_CHAR_NARROW
                    | NB_CHAR_EM
                    | NB_CHAR
                    => output.push_str(&format!("{}",
                                                match c {
                                                    NB_CHAR => "#ffff66",
                                                    NB_CHAR_NARROW => "#9999ff",
                                                    NB_CHAR_EM => "#ff9999",
                                                    _ => unreachable!()
                                                },
                                                c)),
                _ => output.push(c),
            }
        }
        Cow::Owned(output)
    } else {
        input.into()
    }
}

/// Escape characters `<`, `>`, and `&`
///
/// # Examples
///
/// ```
/// use crowbook::escape::escape_html;
/// let s = escape_html("<foo> & <bar>");
/// assert_eq!(&s, "<foo> & <bar>");
/// ```
pub fn escape_html<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    let input = input.into();
    let first = input.chars().position(|c| match c {
        '<'| '>'| '&' => true,
        _ => false
    });
    if let Some(first) = first {
        let mut chars = input.chars().collect::<Vec<_>>();
        let mut output = String::with_capacity(chars.len());
        let rest = chars.split_off(first);
        output.push_str(&chars.into_iter().collect::<String>());
        for c in rest {
            match c {
                '<' => output.push_str("<"),
                '>' => output.push_str(">"),
                '&' => output.push_str("&"),
                _ => output.push(c),
            }
        }

        Cow::Owned(output)
    } else {
        input.into()
    }
}

/// Escape characters for LaTeX
///
/// # Examples
///
/// ```
/// use crowbook::escape::escape_tex;
/// let s = escape_tex("command --foo # calls command with option foo");
/// assert_eq!(&s, r"command -{}-foo \# calls command with option foo");
/// ```
pub fn escape_tex<'a, S: Into<Cow<'a, str>>>(input: S) -> Cow<'a, str> {
    let input = input.into();
    if input.contains(|c| match c {
        '&'|'%'|'$'|'#'|'_'|'{'|'}'|'~'|'^'|'\\'|'-' => true,
        _ => false
    }) {
        let mut output = String::with_capacity(input.len());
        let mut chars:Vec<char> = input.chars().collect();
        chars.push(' '); // add a dummy char for call to .windows()
        // for &[c, next] in chars.windows(2) { // still experimental, uncomment when stable
        for win in chars.windows(2) {
            let c = win[0];
            let next = win[1];
            match c {
                '-' => {
                    if next == '-' {
                        output.push_str(r"-{}"); // if next char is also a -, to avoid tex ligatures
                    } else {
                        output.push(c);
                    }
                },
                '&' => output.push_str(r"\&"),
                '%' => output.push_str(r"\%"),
                '$' => output.push_str(r"\$"),
                '#' => output.push_str(r"\#"),
                '_' => output.push_str(r"\_"),
                '{' => output.push_str(r"\{"),
                '}' => output.push_str(r"\}"),
                '~' => output.push_str(r"\textasciitilde{}" ),
                '^' => output.push_str(r"\textasciicircum{}"),
                '\\' => output.push_str(r"\textbackslash{}"),
                _  => output.push(c)
            }
        }
        Cow::Owned(output)
    } else {
        input
    }
}
<|endoftext|>


<commit_before>extern crate rusty_raft;
extern crate rand;
extern crate capnp;
mod relay_server;
mod mock_state_machine;

use mock_state_machine::*;
use relay_server::*;
use rusty_raft::server::*;
use rusty_raft::rpc::client::*;
use rusty_raft::rpc::RpcError;
use rusty_raft::client_request;
use rusty_raft::client_request::Op;
use rusty_raft::client_request_reply;

use rand::{thread_rng, Rng};
use capnp::message::{Reader};
use capnp::serialize::OwnedSegments;

use std::collections::HashMap;
use std::net::{SocketAddr, ToSocketAddrs};
use std::sync::mpsc::{Receiver, channel};
use std::thread;
use std::time::Duration;

// Container around a MockStateMachine Receiver and a handle
// to the raft server that is running that instance of the state machine
struct StateMachineHandle {
    rx: Receiver<Vec<u8>>,
    server_handle: ServerHandle
}

/// Starts a relay server that binds to num_servers os-assigned local ports
fn start_relay_server (num_servers: u64) -> (RelayServer, HashMap<u64, SocketAddr>) {
    let relay_server = RelayServer::new_with_random_addresses(num_servers);
    let addrs: HashMap<u64, SocketAddr> =
        relay_server.get_bound_addresses()
        .into_iter()
        .enumerate()
        .map(|(i, addr)| (i as u64, addr))
        .collect();

    (relay_server, addrs)
}

/// Starts up a raft server behind each port in the relay server and sets up the mapping
/// from the relay address to the raft server address.
/// Returns a StateMachineHandle that can be used to communicate with the raft server (through the
/// ServerHandle) or check on the state of the MockStateMachine
fn start_raft_servers(relay_server: &mut RelayServer, addrs: &HashMap<u64, SocketAddr>) -> Vec<StateMachineHandle> {
    const HEARTBEAT_TIMEOUT: u64 = 75;
    println!("Starting {} raft servers", addrs.len());

    (0..addrs.len() as u64)
    .map(|i| {
        // create a config object
        let config = Config::new (addrs.clone(), i, "127.0.0.1:0".to_socket_addrs().unwrap().next().unwrap(),
                Duration::from_millis(HEARTBEAT_TIMEOUT));
        let (tx, rx) = channel();
        let state_machine = Box::new(MockStateMachine::new_with_sender(tx));
        let server_handle = start_server(config, move || state_machine).unwrap();

        // map this server through the relay server
        relay_server.relay_address(*addrs.get(&i).unwrap(), server_handle.get_local_addr());
        StateMachineHandle {rx: rx, server_handle: server_handle}
    })
    .collect()
}

#[test]
/// This test is as much a sanity check on our testing code as on the raft code.
/// It just makes sure we can start up a static cluster of servers without crashing
fn it_starts_up_a_cluster() {
    const NUM_SERVERS: u64 = 1;

    {
        let (mut relay_server, addrs) = start_relay_server(NUM_SERVERS);
        let state_machines = start_raft_servers(&mut relay_server, &addrs);
    }
    thread::sleep_ms(500);

    // TODO(jason): Shut down the servers
}

/// Returns a new Rpc
fn create_client_request(op: Op, data: &[u8]) -> Rpc {
    let mut rpc = Rpc::new(CLIENT_REQUEST_OPCODE);
    {
        let mut param_builder = rpc.get_param_builder().init_as::<client_request::Builder>();
        param_builder.set_op(op);
        param_builder.set_data(data);
    }
    rpc
}

#[test]
/// Simple normal case test that starts up a static cluster, sends an entry,
/// and ensures that entry is recieved by all state machines
fn it_replicates_an_entry() {
    const NUM_SERVERS: u64 = 5;
    const REPLICATE_TIMEOUT: u64 = 100;
    const DATA_LENGTH: usize = 1;
    let replicate_timeout = Duration::from_millis(REPLICATE_TIMEOUT);

    let (mut relay_server, addrs) = start_relay_server(NUM_SERVERS);
    let state_machines = start_raft_servers(&mut relay_server, &addrs);

    // TODO: We need a better way of sleeping until a leader is elected
    // Currently we just allow for 3 rounds of split votes...
    thread::sleep_ms(1000);

    // generate the client append RPC and send it to each server
    // only one should respond that they're the leader
    let data: String = thread_rng().gen_ascii_chars().take(DATA_LENGTH).collect();
    let leader_reply_iter = addrs.iter()
        .map(|(_, &to_addr)| {
            let rpc = create_client_request(Op::Write, data.as_bytes());
            rpc.send(to_addr)
        })
        .collect::<Result<Vec<Reader<OwnedSegments>>, RpcError>>()
        .unwrap()
        .into_iter()
        .filter(|msg| {
            let reply = Rpc::get_result_reader(&msg).unwrap();
            reply.get_as::<client_request_reply::Reader>().unwrap().get_success()
        });
    // only 1 server should reply as leader
    assert_eq!(leader_reply_iter.count(), 1);

    // ensure that all state machines replicated the entry
    assert!(state_machines
                .iter()
                .map(|handle| handle.rx.recv_timeout(replicate_timeout).unwrap())
                .all(|vec| vec.iter()
                              .zip(data.clone().into_bytes())
                              .all(|(a, b)| *a == b)));

    // TODO(jason): Shut down the servers
}
<commit_msg>test: Increase test cluster sizes<commit_after>extern crate rusty_raft;
extern crate rand;
extern crate capnp;
mod relay_server;
mod mock_state_machine;

use mock_state_machine::*;
use relay_server::*;
use rusty_raft::server::*;
use rusty_raft::rpc::client::*;
use rusty_raft::rpc::RpcError;
use rusty_raft::client_request;
use rusty_raft::client_request::Op;
use rusty_raft::client_request_reply;

use rand::{thread_rng, Rng};
use capnp::message::{Reader};
use capnp::serialize::OwnedSegments;

use std::collections::HashMap;
use std::net::{SocketAddr, ToSocketAddrs};
use std::sync::mpsc::{Receiver, channel};
use std::thread;
use std::time::Duration;

// Container around a MockStateMachine Receiver and a handle
// to the raft server that is running that instance of the state machine
struct StateMachineHandle {
    rx: Receiver<Vec<u8>>,
    server_handle: ServerHandle
}

/// Starts a relay server that binds to num_servers os-assigned local ports
fn start_relay_server (num_servers: u64) -> (RelayServer, HashMap<u64, SocketAddr>) {
    let relay_server = RelayServer::new_with_random_addresses(num_servers);
    let addrs: HashMap<u64, SocketAddr> =
        relay_server.get_bound_addresses()
        .into_iter()
        .enumerate()
        .map(|(i, addr)| (i as u64, addr))
        .collect();

    (relay_server, addrs)
}

/// Starts up a raft server behind each port in the relay server and sets up the mapping
/// from the relay address to the raft server address.
/// Returns a StateMachineHandle that can be used to communicate with the raft server (through the
/// ServerHandle) or check on the state of the MockStateMachine
fn start_raft_servers(relay_server: &mut RelayServer, addrs: &HashMap<u64, SocketAddr>) -> Vec<StateMachineHandle> {
    const HEARTBEAT_TIMEOUT: u64 = 75;
    println!("Starting {} raft servers", addrs.len());

    (0..addrs.len() as u64)
    .map(|i| {
        // create a config object
        let config = Config::new (addrs.clone(), i, "127.0.0.1:0".to_socket_addrs().unwrap().next().unwrap(),
                Duration::from_millis(HEARTBEAT_TIMEOUT));
        let (tx, rx) = channel();
        let state_machine = Box::new(MockStateMachine::new_with_sender(tx));
        let server_handle = start_server(config, move || state_machine).unwrap();

        // map this server through the relay server
        relay_server.relay_address(*addrs.get(&i).unwrap(), server_handle.get_local_addr());
        StateMachineHandle {rx: rx, server_handle: server_handle}
    })
    .collect()
}

#[test]
/// This test is as much a sanity check on our testing code as on the raft code.
/// It just makes sure we can start up a static cluster of servers without crashing
fn it_starts_up_a_cluster() {
    const NUM_SERVERS: u64 = 3;

    {
        let (mut relay_server, addrs) = start_relay_server(NUM_SERVERS);
        let state_machines = start_raft_servers(&mut relay_server, &addrs);
    }
    thread::sleep_ms(500);

    // TODO(jason): Shut down the servers
}

/// Returns a new Rpc
fn create_client_request(op: Op, data: &[u8]) -> Rpc {
    let mut rpc = Rpc::new(CLIENT_REQUEST_OPCODE);
    {
        let mut param_builder = rpc.get_param_builder().init_as::<client_request::Builder>();
        param_builder.set_op(op);
        param_builder.set_data(data);
    }
    rpc
}

#[test]
/// Simple normal case test that starts up a static cluster, sends an entry,
/// and ensures that entry is recieved by all state machines
fn it_replicates_an_entry() {
    const NUM_SERVERS: u64 = 8;
    const REPLICATE_TIMEOUT: u64 = 100;
    const DATA_LENGTH: usize = 1;
    let replicate_timeout = Duration::from_millis(REPLICATE_TIMEOUT);

    let (mut relay_server, addrs) = start_relay_server(NUM_SERVERS);
    let state_machines = start_raft_servers(&mut relay_server, &addrs);

    // TODO: We need a better way of sleeping until a leader is elected
    // Currently we just allow for 3 rounds of split votes...
    thread::sleep_ms(1000);

    // generate the client append RPC and send it to each server
    // only one should respond that they're the leader
    let data: String = thread_rng().gen_ascii_chars().take(DATA_LENGTH).collect();
    let leader_reply_iter = addrs.iter()
        .map(|(_, &to_addr)| {
            let rpc = create_client_request(Op::Write, data.as_bytes());
            rpc.send(to_addr)
        })
        .collect::<Result<Vec<Reader<OwnedSegments>>, RpcError>>()
        .unwrap()
        .into_iter()
        .filter(|msg| {
            let reply = Rpc::get_result_reader(&msg).unwrap();
            reply.get_as::<client_request_reply::Reader>().unwrap().get_success()
        });
    // only 1 server should reply as leader
    assert_eq!(leader_reply_iter.count(), 1);

    // ensure that all state machines replicated the entry
    assert!(state_machines
                .iter()
                .map(|handle| handle.rx.recv_timeout(replicate_timeout).unwrap())
                .all(|vec| vec.iter()
                              .zip(data.clone().into_bytes())
                              .all(|(a, b)| *a == b)));

    // TODO(jason): Shut down the servers
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::Thread;
use std::cmp;

use messages::*;
use teller::*;

struct Market {
  id: usize,
  tellers: HashMap<usize, Sender<TellerMessages>>,
  actors: HashMap<usize, Sender<ActorMessages>>,
  active_transactions: Vec<(TransactionRequest, TransactionRequest)>,
  pending_transactions: Vec<(TransactionRequest, TransactionRequest)>,
  committed_actors: Vec<usize>,
}

//Called on a new thread
pub fn start_market(market_id: usize, market_tx: Sender<MarketMessages>, market_rx: Receiver<MarketMessages>) {
  //Create Market struct
  let mut market = Market {id:market_id,
                             tellers: HashMap::new(),
                             actors: HashMap::new(),
                             active_transactions: vec![],
                             committed_actors: vec![],
                             pending_transactions: vec![]};
  //TODO: Initialize Tellers
  for i in 0..1 {
    let (tx, rx): (Sender<TellerMessages>, Receiver<TellerMessages>) = channel();
    market.tellers.insert(i, tx);
    let market_tx_clone = market_tx.clone();
    Thread::spawn(move ||
      {start_teller(i, market_tx_clone, rx);});
  }

  //Start the receive loop
  loop {
    let message = market_rx.recv().unwrap();
    match message {
        MarketMessages::SellRequest(request) => {route(false, request, &market)},
        MarketMessages::BuyRequest(request) => {route(true, request, &market)},
        MarketMessages::Commit(actor_id) => {
          println!("Commit {}", actor_id);
          if has_active_transaction(&market, actor_id) {
            market.committed_actors.push(actor_id);
            let tup = get_active_transaction_involving(&market, actor_id);
            if contains(&market.committed_actors, tup.0.actor_id) && contains(&market.committed_actors, tup.1.actor_id) {
              //make the transaction complete. Tell the actors to commit;
              remove(&mut market.committed_actors, tup.0.actor_id);
              remove(&mut market.committed_actors, tup.1.actor_id);

              route_actor_message(&market, tup.0.actor_id, ActorMessages::CommitTransaction(tup.1.clone()));
              route_actor_message(&market, tup.1.actor_id, ActorMessages::CommitTransaction(tup.0.clone()));

              remove_active_transaction(&mut market, tup);
              //TODO: Move Pending Transaction to Active and notify the actors.
            }
          }
        }
        MarketMessages::Cancel(actor_id) => {
          println!("Cancel {}", actor_id);
          if has_active_transaction(&market, actor_id) {
            let tup = get_active_transaction_involving(&market, actor_id);
            route_actor_message(&market, tup.0.actor_id, ActorMessages::AbortTransaction);
            route_actor_message(&market, tup.1.actor_id, ActorMessages::AbortTransaction);

            //remove it from active
            remove_active_transaction(&mut market, tup);

            //TODO: Move Pending Transaction to Active and notify for the actors that we just
            //aborted.
          }
          },
        MarketMessages::RegisterActor(actor_id, actor_tx) => {
          let temp_clone = actor_tx.clone();
          market.actors.insert(actor_id, actor_tx);},
        MarketMessages::MatchRequest(mut buyer, mut seller) => {
          if has_active_transaction(&market, buyer.actor_id) || has_active_transaction(&market, seller.actor_id) {
            //add to pending transactions
            market.pending_transactions.push((buyer, seller));
          }
          else {
            //add to active transactions and notify both.
            buyer.price = seller.price;
            let smaller_quantity = cmp::min(seller.quantity, buyer.quantity);
            buyer.quantity = smaller_quantity;
            seller.quantity = smaller_quantity;
            let amount_to_pay = buyer.quantity * buyer.price;
            let buyer_request = MoneyRequest {market_id: market.id, amount: amount_to_pay};
            let seller_request = StockRequest {market_id: market.id, stock_id: seller.stock_id, quantity: seller.quantity};

            route_actor_message(&market, buyer.actor_id, ActorMessages::MoneyRequest(buyer_request));
            route_actor_message(&market, seller.actor_id, ActorMessages::StockRequest(seller_request));
            market.active_transactions.push((buyer, seller));
          }
          },
    }
  }
}

fn remove<T:PartialEq>(vec: &mut Vec<T>, element: T) {
  for i in 0..vec.len() {
    if vec[i] == element {
      vec.remove(i);
      break;
    }
  }
}

fn contains<T:PartialEq>(vec: &Vec<T>, element: T) -> bool {
  for test_element in vec.iter() {
    if element == *test_element {
      return true;
    }
  }
  false
}

fn remove_active_transaction(market: &mut Market, tup: (TransactionRequest, TransactionRequest)) {
  for i in 0..market.active_transactions.len() {
    let test_active = market.active_transactions[i].clone();
    if  test_active.0 == tup.0 && test_active.1 == tup.1 {
       market.active_transactions.remove(i);
       break;
    }
  }
}

fn get_active_transaction_involving(market_ref : &Market, id: usize) -> (TransactionRequest, TransactionRequest) {
  for transaction_pair in market_ref.active_transactions.iter() {
    if transaction_pair.0.actor_id == id || transaction_pair.1.actor_id == id {
      return transaction_pair.clone();
    }
  }
  (TransactionRequest {transaction_id: 0, actor_id: 0, stock_id: 0, price: 0, quantity: 0}, TransactionRequest {transaction_id: 0, actor_id: 0, stock_id: 0, price: 0, quantity: 0})
}

fn has_active_transaction(market_ref : &Market, id: usize) -> bool {
  for transaction_pair in market_ref.active_transactions.iter() {
    if transaction_pair.0.actor_id == id || transaction_pair.1.actor_id == id {
      return true;
    }
  }
  false
}

fn route_actor_message(market: & Market, actor_id: usize, message: ActorMessages) {
  match market.actors.get(&actor_id) {
    Some(channel) => {channel.send(message);},
    None => {}
  }
}

fn route(buying: bool, transaction: TransactionRequest, market: & Market) {
  let tx;
  match market.tellers.get(&transaction.stock_id) {
    Some(channel) => {tx = channel;},
    None => {return;}
  }
  if buying {
    tx.send(TellerMessages::BuyRequest(transaction));
  }
  else {
    tx.send(TellerMessages::SellRequest(transaction));
  };
}
<commit_msg>Market should now move pending into active. Still untested until we get an actor wanting to sell and an actor wanting to buy.<commit_after>use std::collections::HashMap;
use std::sync::mpsc::{Sender, Receiver, channel};
use std::thread::Thread;
use std::cmp;

use messages::*;
use teller::*;

struct Market {
  id: usize,
  tellers: HashMap<usize, Sender<TellerMessages>>,
  actors: HashMap<usize, Sender<ActorMessages>>,
  active_transactions: Vec<(TransactionRequest, TransactionRequest)>,
  pending_transactions: Vec<(TransactionRequest, TransactionRequest)>,
  committed_actors: Vec<usize>,
}

//Called on a new thread
pub fn start_market(market_id: usize, market_tx: Sender<MarketMessages>, market_rx: Receiver<MarketMessages>) {
  //Create Market struct
  let mut market = Market {id:market_id,
                             tellers: HashMap::new(),
                             actors: HashMap::new(),
                             active_transactions: vec![],
                             committed_actors: vec![],
                             pending_transactions: vec![]};
  //TODO: Initialize Tellers
  for i in 0..1 {
    let (tx, rx): (Sender<TellerMessages>, Receiver<TellerMessages>) = channel();
    market.tellers.insert(i, tx);
    let market_tx_clone = market_tx.clone();
    Thread::spawn(move ||
      {start_teller(i, market_tx_clone, rx);});
  }

  //Start the receive loop
  loop {
    let message = market_rx.recv().unwrap();
    match message {
        MarketMessages::SellRequest(request) => {route(false, request, &market)},
        MarketMessages::BuyRequest(request) => {route(true, request, &market)},
        MarketMessages::Commit(actor_id) => {
          println!("Commit {}", actor_id);
          if has_active_transaction(&market, actor_id) {
            market.committed_actors.push(actor_id);
            let tup = get_active_transaction_involving(&market, actor_id);
            if contains(&market.committed_actors, tup.0.actor_id) && contains(&market.committed_actors, tup.1.actor_id) {
              //make the transaction complete. Tell the actors to commit;
              remove(&mut market.committed_actors, tup.0.actor_id);
              remove(&mut market.committed_actors, tup.1.actor_id);

              route_actor_message(&market, tup.0.actor_id, ActorMessages::CommitTransaction(tup.1.clone()));
              route_actor_message(&market, tup.1.actor_id, ActorMessages::CommitTransaction(tup.0.clone()));

              remove_active_transaction(&mut market, &tup);
              move_pending_to_active(&mut market, tup.0.actor_id, tup.1.actor_id);
            }
          }
        }
        MarketMessages::Cancel(actor_id) => {
          println!("Cancel {}", actor_id);
          if has_active_transaction(&market, actor_id) {
            let tup = get_active_transaction_involving(&market, actor_id);
            route_actor_message(&market, tup.0.actor_id, ActorMessages::AbortTransaction);
            route_actor_message(&market, tup.1.actor_id, ActorMessages::AbortTransaction);

            //remove it from active
            remove_active_transaction(&mut market, &tup);

            move_pending_to_active(&mut market, tup.0.actor_id, tup.1.actor_id);
          }
          },
        MarketMessages::RegisterActor(actor_id, actor_tx) => {
          let temp_clone = actor_tx.clone();
          market.actors.insert(actor_id, actor_tx);},
        MarketMessages::MatchRequest(mut buyer, mut seller) => {
          if has_active_transaction(&market, buyer.actor_id) || has_active_transaction(&market, seller.actor_id) {
            //add to pending transactions
            market.pending_transactions.push((buyer, seller));
          }
          else {
            activate_transactions(&mut market, buyer, seller);
          }
          },
    }
  }
}

fn activate_transactions(market: &mut Market, mut buyer: TransactionRequest, mut seller: TransactionRequest) {
  //add to active transactions and notify both.
  buyer.price = seller.price;
  let smaller_quantity = cmp::min(seller.quantity, buyer.quantity);
  buyer.quantity = smaller_quantity;
  seller.quantity = smaller_quantity;
  let amount_to_pay = buyer.quantity * buyer.price;
  let buyer_request = MoneyRequest {market_id: market.id, amount: amount_to_pay};
  let seller_request = StockRequest {market_id: market.id, stock_id: seller.stock_id, quantity: seller.quantity};

  route_actor_message(&market, buyer.actor_id, ActorMessages::MoneyRequest(buyer_request));
  route_actor_message(&market, seller.actor_id, ActorMessages::StockRequest(seller_request));
  market.active_transactions.push((buyer, seller));
}

fn move_pending_to_active(market: &mut Market, actor1: usize, actor2: usize) {
  //actor 1
  for i in 0..market.pending_transactions.len() {
    let pending_transaction = market.pending_transactions[i].clone();
    if (pending_transaction.0.actor_id == actor1) {
      if !has_active_transaction(&market, pending_transaction.1.actor_id) {
        activate_transactions(market, pending_transaction.0.clone(), pending_transaction.1.clone());
        remove(&mut market.pending_transactions, pending_transaction);
        break;
      }
    }
    if (pending_transaction.1.actor_id == actor1) {
      if !has_active_transaction(&market, pending_transaction.0.actor_id) {
        activate_transactions(market, pending_transaction.0.clone(), pending_transaction.1.clone());
        remove(&mut market.pending_transactions, pending_transaction);
        break;
      }
    }
  }

  //actor 2
  for i in 0..market.pending_transactions.len() {
    let pending_transaction = market.pending_transactions[i].clone();
    if (pending_transaction.0.actor_id == actor2) {
      if !has_active_transaction(&market, pending_transaction.1.actor_id) {
        activate_transactions(market, pending_transaction.0.clone(), pending_transaction.1.clone());
        remove(&mut market.pending_transactions, pending_transaction);
        break;
      }
    }
    if (pending_transaction.1.actor_id == actor2) {
      if !has_active_transaction(&market, pending_transaction.0.actor_id) {
        activate_transactions(market, pending_transaction.0.clone(), pending_transaction.1.clone());
        remove(&mut market.pending_transactions, pending_transaction);
        break;
      }
    }
  }
}

fn remove<T:PartialEq>(vec: &mut Vec<T>, element: T) {
  for i in 0..vec.len() {
    if vec[i] == element {
      vec.remove(i);
      break;
    }
  }
}

fn contains<T:PartialEq>(vec: &Vec<T>, element: T) -> bool {
  for test_element in vec.iter() {
    if element == *test_element {
      return true;
    }
  }
  false
}

fn remove_active_transaction(market: &mut Market, tup: &(TransactionRequest, TransactionRequest)) {
  for i in 0..market.active_transactions.len() {
    let test_active = market.active_transactions[i].clone();
    if  test_active.0 == tup.0 && test_active.1 == tup.1 {
       market.active_transactions.remove(i);
       break;
    }
  }
}

fn get_active_transaction_involving(market_ref : &Market, id: usize) -> (TransactionRequest, TransactionRequest) {
  for transaction_pair in market_ref.active_transactions.iter() {
    if transaction_pair.0.actor_id == id || transaction_pair.1.actor_id == id {
      return transaction_pair.clone();
    }
  }
  (TransactionRequest {transaction_id: 0, actor_id: 0, stock_id: 0, price: 0, quantity: 0}, TransactionRequest {transaction_id: 0, actor_id: 0, stock_id: 0, price: 0, quantity: 0})
}

fn has_active_transaction(market_ref : &Market, id: usize) -> bool {
  for transaction_pair in market_ref.active_transactions.iter() {
    if transaction_pair.0.actor_id == id || transaction_pair.1.actor_id == id {
      return true;
    }
  }
  false
}

fn route_actor_message(market: & Market, actor_id: usize, message: ActorMessages) {
  match market.actors.get(&actor_id) {
    Some(channel) => {channel.send(message);},
    None => {}
  }
}

fn route(buying: bool, transaction: TransactionRequest, market: & Market) {
  let tx;
  match market.tellers.get(&transaction.stock_id) {
    Some(channel) => {tx = channel;},
    None => {return;}
  }
  if buying {
    tx.send(TellerMessages::BuyRequest(transaction));
  }
  else {
    tx.send(TellerMessages::SellRequest(transaction));
  };
}
<|endoftext|>