<commit_before>extern crate rocksdb;
extern crate rustc_serialize;

use std::collections::HashMap;

use self::rustc_serialize::json::{JsonEvent, Parser, StackElement};

use error::Error;
use key_builder::KeyBuilder;
use records_capnp::payload;
use stems::Stems;

// 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>;

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

pub trait Indexable {

}

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

impl Shredder {
    pub fn new() -> Shredder {
        Shredder{
            kb: KeyBuilder::new(),
            ignore_children: 0,
            doc_id: String::new(),
        }
    }

    fn add_entries(&mut self, text: &String, docseq: u64, batch: &mut rocksdb::WriteBatch) ->
            Result<(), Error> {
        let stems = Stems::new(text.as_str());
        let mut word_to_word_infos = HashMap::new();

        for stem in stems {
            let word_infos = word_to_word_infos.entry(stem.stemmed).or_insert(Vec::new());
            word_infos.push(WordInfo{
                stemmed_offset: stem.stemmed_offset as u64,
                suffix_text: stem.suffix.to_string(),
                suffix_offset: stem.suffix_offset as u64,
            });
        }
        for (stemmed, word_infos) in word_to_word_infos {
            let mut message = ::capnp::message::Builder::new_default();
            {
                let capn_payload = message.init_root::<payload::Builder>();
                let mut capn_wordinfos = capn_payload.init_wordinfos(word_infos.len() as u32);
                for (pos, word_info) in word_infos.iter().enumerate() {
                    let mut capn_wordinfo = capn_wordinfos.borrow().get(pos as u32);
                    capn_wordinfo.set_stemmed_offset(word_info.stemmed_offset);
                    capn_wordinfo.set_suffix_text(&word_info.suffix_text);
                    capn_wordinfo.set_suffix_offset(word_info.suffix_offset);
                }
            }

            let mut bytes = Vec::new();
            ::capnp::serialize_packed::write_message(&mut bytes, &message).unwrap();
            let key = self.kb.stemmed_word_key(&stemmed, docseq);
            try!(batch.put(&key.into_bytes(), &bytes));

        }
        let key = self.kb.value_key(docseq);
        try!(batch.put(&key.into_bytes(), &text.as_bytes()));

        Ok(())
    }

    // 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.kb.last_pushed_keypath_is_object_key() {
            match stack_element {
                Some(StackElement::Key(key)) => {
                    if self.kb.keypath_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.kb.last_pushed_keypath_is_object_key() {
            if let Some(StackElement::Key(key)) = stack_element {
                if self.kb.keypath_segments_len() == 1 && 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.kb.pop_object_key();
                    self.kb.push_object_key(key);
                }
            }
        }
        Ok(())
    }

    pub fn shred(&mut self, json: &str, docseq: u64, batch: &mut rocksdb::WriteBatch) ->
            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.kb.push_object_key("");
                    }
                },
                Some(JsonEvent::ObjectEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {
                        self.kb.pop_object_key();
                        self.kb.inc_top_array_offset();
                    }
                },
                Some(JsonEvent::ArrayStart) => {
                    try!(self.maybe_push_key(parser.stack().top()));
                    if self.ignore_children > 0 {
                        self.ignore_children += 1;
                    } else {
                        self.kb.push_array();
                    }
                },
                Some(JsonEvent::ArrayEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {
                        self.kb.pop_array();
                        self.kb.inc_top_array_offset();
                    }
                },
                Some(JsonEvent::StringValue(value)) => {
                    // No children to ignore
                    if self.ignore_children == 0 {
                        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.kb.pop_object_key();
                                self.kb.push_object_key(&key);

                                try!(self.add_entries(&value, docseq, batch));
                                self.kb.inc_top_array_offset();
                            },
                            ObjectKeyTypes::NoKey => {
                                try!(self.add_entries(&value, docseq, batch));
                                self.kb.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;
            }
        }
        Ok(self.doc_id.clone())
   }
}

#[cfg(test)]
mod tests {
    extern crate rocksdb;
    use std::str;
    use records_capnp;
    use super::{WordInfo};

    fn wordinfos_from_rocks(rocks: rocksdb::DB) -> Vec<(String, Vec<WordInfo>)> {
        let mut result = Vec::new();
        for (key, value) in rocks.iterator(rocksdb::IteratorMode::Start) {
            let mut ref_value = &*value;
            let message_reader = ::capnp::serialize_packed::read_message(
                &mut ref_value, ::capnp::message::ReaderOptions::new()).unwrap();
            let payload = message_reader.get_root::<records_capnp::payload::Reader>().unwrap();

            let mut wordinfos = Vec::new();
            for wi in payload.get_wordinfos().unwrap().iter() {
                wordinfos.push(WordInfo{
                    stemmed_offset: wi.get_stemmed_offset(),
                    suffix_text: wi.get_suffix_text().unwrap().to_string(),
                    suffix_offset: wi.get_suffix_offset(),
                });
            }
            let key_string = unsafe { str::from_utf8_unchecked((&key)) }.to_string();
            result.push((key_string, wordinfos));
        }
        result
    }

    #[test]
    fn test_shred_nested() {
        let mut shredder = super::Shredder::new();
        let json = r#"{"some": ["array", "data", ["also", "nested"]]}"#;
        let docseq = 123;
        let batch = rocksdb::WriteBatch::default();
        shredder.shred(json, docseq, &batch).unwrap();

        let rocks = rocksdb::DB::open_default("target/tests/test_shred_netsted").unwrap();
        rocks.write(batch).unwrap();
        let result = wordinfos_from_rocks(rocks);

        let expected = vec![
            ("W.some$!array#123,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }]),
            ("W.some$!data#123,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }]),
            ("W.some$$!also#123,2,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }]),
            ("W.some$$!nest#123,2,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "ed".to_string(), suffix_offset: 4 }]),
            ];
        assert_eq!(result, expected);
    }

    #[test]
    // NOTE vmx 2016-12-06: This test is intentionally made to fail (hence ignored) as the current
    // current tokenizer does the wrong thing when it comes to numbers within words. It's left
    // here as a reminder to fix that
    #[ignore]
    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;
        let batch = rocksdb::WriteBatch::default();
        shredder.shred(json, docseq, &batch).unwrap();

        let rocks = rocksdb::DB::open_default("target/tests/test_shred_objects").unwrap();
        rocks.write(batch).unwrap();
        let result = wordinfos_from_rocks(rocks);
        println!("result: {:?}", result);
        let expected = vec![
            ("W.A$.B!b1#1234,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }]),
            ("W.A$.B!b2vmx#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "B2 VMX ".to_string(),
                           suffix_offset: 0 }]),
            ("W.A$.B!three#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 10, suffix_text: "".to_string(), suffix_offset: 15 }]),
            ("W.A$.B!two#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 6, suffix_text: " ".to_string(), suffix_offset: 9 }]),
            ("W.A$.C!..#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }]),
            ("W.A$.C!..#1234,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }]),
            ("W.A$.C!c2#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 2, suffix_text: "C2".to_string(), suffix_offset: 2 }]),
            ("W.A$.C!c2#1234,1".to_string(), vec![
                WordInfo { stemmed_offset: 2, suffix_text: "C2".to_string(), suffix_offset: 2 }]),
            ];
        assert_eq!(result, expected);
    }

    #[test]
    fn test_shred_empty_object() {
        let mut shredder = super::Shredder::new();
        let json = r#"{}"#;
        let docseq = 123;
        let batch = rocksdb::WriteBatch::default();
        shredder.shred(json, docseq, &batch).unwrap();

        let rocks = rocksdb::DB::open_default("target/tests/test_shred_empty_object").unwrap();
        rocks.write(batch).unwrap();
        let result = wordinfos_from_rocks(rocks);
        assert!(result.is_empty());
    }
}
<commit_msg>Fix broken tests by making mut ref<commit_after>extern crate rocksdb;
extern crate rustc_serialize;

use std::collections::HashMap;

use self::rustc_serialize::json::{JsonEvent, Parser, StackElement};

use error::Error;
use key_builder::KeyBuilder;
use records_capnp::payload;
use stems::Stems;

// 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>;

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

pub trait Indexable {

}

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

impl Shredder {
    pub fn new() -> Shredder {
        Shredder{
            kb: KeyBuilder::new(),
            ignore_children: 0,
            doc_id: String::new(),
        }
    }

    fn add_entries(&mut self, text: &String, docseq: u64, batch: &mut rocksdb::WriteBatch) ->
            Result<(), Error> {
        let stems = Stems::new(text.as_str());
        let mut word_to_word_infos = HashMap::new();

        for stem in stems {
            let word_infos = word_to_word_infos.entry(stem.stemmed).or_insert(Vec::new());
            word_infos.push(WordInfo{
                stemmed_offset: stem.stemmed_offset as u64,
                suffix_text: stem.suffix.to_string(),
                suffix_offset: stem.suffix_offset as u64,
            });
        }
        for (stemmed, word_infos) in word_to_word_infos {
            let mut message = ::capnp::message::Builder::new_default();
            {
                let capn_payload = message.init_root::<payload::Builder>();
                let mut capn_wordinfos = capn_payload.init_wordinfos(word_infos.len() as u32);
                for (pos, word_info) in word_infos.iter().enumerate() {
                    let mut capn_wordinfo = capn_wordinfos.borrow().get(pos as u32);
                    capn_wordinfo.set_stemmed_offset(word_info.stemmed_offset);
                    capn_wordinfo.set_suffix_text(&word_info.suffix_text);
                    capn_wordinfo.set_suffix_offset(word_info.suffix_offset);
                }
            }

            let mut bytes = Vec::new();
            ::capnp::serialize_packed::write_message(&mut bytes, &message).unwrap();
            let key = self.kb.stemmed_word_key(&stemmed, docseq);
            try!(batch.put(&key.into_bytes(), &bytes));

        }
        let key = self.kb.value_key(docseq);
        try!(batch.put(&key.into_bytes(), &text.as_bytes()));

        Ok(())
    }

    // 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.kb.last_pushed_keypath_is_object_key() {
            match stack_element {
                Some(StackElement::Key(key)) => {
                    if self.kb.keypath_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.kb.last_pushed_keypath_is_object_key() {
            if let Some(StackElement::Key(key)) = stack_element {
                if self.kb.keypath_segments_len() == 1 && 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.kb.pop_object_key();
                    self.kb.push_object_key(key);
                }
            }
        }
        Ok(())
    }

    pub fn shred(&mut self, json: &str, docseq: u64, batch: &mut rocksdb::WriteBatch) ->
            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.kb.push_object_key("");
                    }
                },
                Some(JsonEvent::ObjectEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {
                        self.kb.pop_object_key();
                        self.kb.inc_top_array_offset();
                    }
                },
                Some(JsonEvent::ArrayStart) => {
                    try!(self.maybe_push_key(parser.stack().top()));
                    if self.ignore_children > 0 {
                        self.ignore_children += 1;
                    } else {
                        self.kb.push_array();
                    }
                },
                Some(JsonEvent::ArrayEnd) => {
                    if self.ignore_children > 0 {
                        self.ignore_children -= 1;
                    } else {
                        self.kb.pop_array();
                        self.kb.inc_top_array_offset();
                    }
                },
                Some(JsonEvent::StringValue(value)) => {
                    // No children to ignore
                    if self.ignore_children == 0 {
                        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.kb.pop_object_key();
                                self.kb.push_object_key(&key);

                                try!(self.add_entries(&value, docseq, batch));
                                self.kb.inc_top_array_offset();
                            },
                            ObjectKeyTypes::NoKey => {
                                try!(self.add_entries(&value, docseq, batch));
                                self.kb.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;
            }
        }
        Ok(self.doc_id.clone())
   }
}

#[cfg(test)]
mod tests {
    extern crate rocksdb;
    use std::str;
    use records_capnp;
    use super::{WordInfo};

    fn wordinfos_from_rocks(rocks: rocksdb::DB) -> Vec<(String, Vec<WordInfo>)> {
        let mut result = Vec::new();
        for (key, value) in rocks.iterator(rocksdb::IteratorMode::Start) {
            let mut ref_value = &*value;
            let message_reader = ::capnp::serialize_packed::read_message(
                &mut ref_value, ::capnp::message::ReaderOptions::new()).unwrap();
            let payload = message_reader.get_root::<records_capnp::payload::Reader>().unwrap();

            let mut wordinfos = Vec::new();
            for wi in payload.get_wordinfos().unwrap().iter() {
                wordinfos.push(WordInfo{
                    stemmed_offset: wi.get_stemmed_offset(),
                    suffix_text: wi.get_suffix_text().unwrap().to_string(),
                    suffix_offset: wi.get_suffix_offset(),
                });
            }
            let key_string = unsafe { str::from_utf8_unchecked((&key)) }.to_string();
            result.push((key_string, wordinfos));
        }
        result
    }

    #[test]
    fn test_shred_nested() {
        let mut shredder = super::Shredder::new();
        let json = r#"{"some": ["array", "data", ["also", "nested"]]}"#;
        let docseq = 123;
        let batch = rocksdb::WriteBatch::default();
        shredder.shred(json, docseq, &mut batch).unwrap();

        let rocks = rocksdb::DB::open_default("target/tests/test_shred_netsted").unwrap();
        rocks.write(batch).unwrap();
        let result = wordinfos_from_rocks(rocks);

        let expected = vec![
            ("W.some$!array#123,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 5 }]),
            ("W.some$!data#123,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }]),
            ("W.some$$!also#123,2,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 4 }]),
            ("W.some$$!nest#123,2,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "ed".to_string(), suffix_offset: 4 }]),
            ];
        assert_eq!(result, expected);
    }

    #[test]
    // NOTE vmx 2016-12-06: This test is intentionally made to fail (hence ignored) as the current
    // current tokenizer does the wrong thing when it comes to numbers within words. It's left
    // here as a reminder to fix that
    #[ignore]
    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;
        let batch = rocksdb::WriteBatch::default();
        shredder.shred(json, docseq, &mut batch).unwrap();

        let rocks = rocksdb::DB::open_default("target/tests/test_shred_objects").unwrap();
        rocks.write(batch).unwrap();
        let result = wordinfos_from_rocks(rocks);
        println!("result: {:?}", result);
        let expected = vec![
            ("W.A$.B!b1#1234,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }]),
            ("W.A$.B!b2vmx#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "B2 VMX ".to_string(),
                           suffix_offset: 0 }]),
            ("W.A$.B!three#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 10, suffix_text: "".to_string(), suffix_offset: 15 }]),
            ("W.A$.B!two#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 6, suffix_text: " ".to_string(), suffix_offset: 9 }]),
            ("W.A$.C!..#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }]),
            ("W.A$.C!..#1234,1".to_string(), vec![
                WordInfo { stemmed_offset: 0, suffix_text: "".to_string(), suffix_offset: 2 }]),
            ("W.A$.C!c2#1234,0".to_string(), vec![
                WordInfo { stemmed_offset: 2, suffix_text: "C2".to_string(), suffix_offset: 2 }]),
            ("W.A$.C!c2#1234,1".to_string(), vec![
                WordInfo { stemmed_offset: 2, suffix_text: "C2".to_string(), suffix_offset: 2 }]),
            ];
        assert_eq!(result, expected);
    }

    #[test]
    fn test_shred_empty_object() {
        let mut shredder = super::Shredder::new();
        let json = r#"{}"#;
        let docseq = 123;
        let batch = rocksdb::WriteBatch::default();
        shredder.shred(json, docseq, &mut batch).unwrap();

        let rocks = rocksdb::DB::open_default("target/tests/test_shred_empty_object").unwrap();
        rocks.write(batch).unwrap();
        let result = wordinfos_from_rocks(rocks);
        assert!(result.is_empty());
    }
}
<|endoftext|>


<commit_before>use tokenizer::new;
use tokenizer::Tokenizer;
use tokenizer::StateFunction;
use token::Token;
use token::Category;

fn initial_state(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            if lexer.token_position == lexer.token_start {
                let remaining_data = lexer.data
                    .slice_from(lexer.token_position).to_string();

                if remaining_data.starts_with("</") {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(2, Category::Text);
                    return Some(StateFunction(inside_tag))
                }
            }
            match c {
                '<' => {
                    lexer.tokenize_next(1, Category::Text);
                    return Some(StateFunction(start_of_tag));
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Text);
                    lexer.advance();
                    lexer.states.push(StateFunction(initial_state));
                    return Some(StateFunction(whitespace));
                },
                _ => {
                    lexer.advance();
                }
            }

            Some(StateFunction(initial_state))
        }

        None => {
            lexer.tokenize(Category::Text);
            None
        }
    }
}

fn start_of_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(start_of_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    Some(StateFunction(inside_string))
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '=' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::AssignmentOperator);
                    Some(StateFunction(inside_tag))
                }
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    if lexer.token_position == lexer.token_start {
                        let remaining_data = lexer.data
                            .slice_from(lexer.token_position).to_string();

                        if remaining_data.starts_with("/>") {
                            lexer.tokenize(Category::Identifier);
                            lexer.tokenize_next(2, Category::Text);
                            return Some(StateFunction(initial_state))
                        }
                    }

                    lexer.advance();
                    Some(StateFunction(inside_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_string(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.advance();
                    lexer.tokenize(Category::String);
                    Some(StateFunction(inside_tag))
                },
                '\\' => {
                    lexer.advance();
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
            }
        }

        None => {
            lexer.tokenize(Category::String);
            None
        }
    }
}

fn whitespace(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.advance();
                    Some(StateFunction(whitespace))
                },
                _ => {
                    lexer.tokenize(Category::Whitespace);
                    Some(lexer.states.pop().unwrap())
                }
            }
        }

        None => {
            lexer.tokenize(Category::Whitespace);
            None
        }
    }
}

pub fn lex(data: &str) -> Vec<Token> {
    let mut lexer = new(data);
    let mut state_function = StateFunction(initial_state);
    loop {
        let StateFunction(actual_function) = state_function;
        match actual_function(&mut lexer) {
            Some(f) => state_function = f,
            None => return lexer.tokens(),
        }
    }
}

mod tests {
    use super::lex;
    use token::Token;
    use token::Category;
    use std::old_io::{File, Open, Read};

    #[test]
    fn it_works() {
        let data = include_str!("../../test_data/data.xml");
        let tokens = lex(data);
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: "=".to_string(), category: Category::AssignmentOperator },
            Token{ lexeme: "\"value\"".to_string(), category: Category::String },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "self_closing_tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "/>".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_garbage() {
        let tokens = lex("} adwyx123&*_ ");
        let expected_tokens = vec![
            Token{ lexeme: "}".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "adwyx123&*_".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_open_strings() {
        let tokens = lex("<tag \"open!>");
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "\"open!>".to_string(), category: Category::String },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_utf8_data() {
        let tokens = lex("différent");
        let expected_tokens = vec![
            Token{ lexeme: "différent".to_string(), category: Category::Text },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }
}
<commit_msg>Drop unused reference to old IO libraries.<commit_after>use tokenizer::new;
use tokenizer::Tokenizer;
use tokenizer::StateFunction;
use token::Token;
use token::Category;

fn initial_state(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            if lexer.token_position == lexer.token_start {
                let remaining_data = lexer.data
                    .slice_from(lexer.token_position).to_string();

                if remaining_data.starts_with("</") {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(2, Category::Text);
                    return Some(StateFunction(inside_tag))
                }
            }
            match c {
                '<' => {
                    lexer.tokenize_next(1, Category::Text);
                    return Some(StateFunction(start_of_tag));
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Text);
                    lexer.advance();
                    lexer.states.push(StateFunction(initial_state));
                    return Some(StateFunction(whitespace));
                },
                _ => {
                    lexer.advance();
                }
            }

            Some(StateFunction(initial_state))
        }

        None => {
            lexer.tokenize(Category::Text);
            None
        }
    }
}

fn start_of_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(start_of_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    Some(StateFunction(inside_string))
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '=' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::AssignmentOperator);
                    Some(StateFunction(inside_tag))
                }
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    if lexer.token_position == lexer.token_start {
                        let remaining_data = lexer.data
                            .slice_from(lexer.token_position).to_string();

                        if remaining_data.starts_with("/>") {
                            lexer.tokenize(Category::Identifier);
                            lexer.tokenize_next(2, Category::Text);
                            return Some(StateFunction(initial_state))
                        }
                    }

                    lexer.advance();
                    Some(StateFunction(inside_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_string(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.advance();
                    lexer.tokenize(Category::String);
                    Some(StateFunction(inside_tag))
                },
                '\\' => {
                    lexer.advance();
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
            }
        }

        None => {
            lexer.tokenize(Category::String);
            None
        }
    }
}

fn whitespace(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.advance();
                    Some(StateFunction(whitespace))
                },
                _ => {
                    lexer.tokenize(Category::Whitespace);
                    Some(lexer.states.pop().unwrap())
                }
            }
        }

        None => {
            lexer.tokenize(Category::Whitespace);
            None
        }
    }
}

pub fn lex(data: &str) -> Vec<Token> {
    let mut lexer = new(data);
    let mut state_function = StateFunction(initial_state);
    loop {
        let StateFunction(actual_function) = state_function;
        match actual_function(&mut lexer) {
            Some(f) => state_function = f,
            None => return lexer.tokens(),
        }
    }
}

mod tests {
    use super::lex;
    use token::Token;
    use token::Category;

    #[test]
    fn it_works() {
        let data = include_str!("../../test_data/data.xml");
        let tokens = lex(data);
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: "=".to_string(), category: Category::AssignmentOperator },
            Token{ lexeme: "\"value\"".to_string(), category: Category::String },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "self_closing_tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "/>".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_garbage() {
        let tokens = lex("} adwyx123&*_ ");
        let expected_tokens = vec![
            Token{ lexeme: "}".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "adwyx123&*_".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_open_strings() {
        let tokens = lex("<tag \"open!>");
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "\"open!>".to_string(), category: Category::String },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_utf8_data() {
        let tokens = lex("différent");
        let expected_tokens = vec![
            Token{ lexeme: "différent".to_string(), category: Category::Text },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }
}
<|endoftext|>


<commit_before>use {Cursive, Printer, With};
use align::HAlign;
use direction::Direction;
use event::*;
use theme::ColorStyle;
use unicode_width::UnicodeWidthStr;
use vec::Vec2;
use view::View;

/// Simple text label with a callback when <Enter> is pressed.
///
/// A button shows its content in a single line and has a fixed size.
///
/// # Examples
///
/// ```
/// # use cursive::views::Button;
/// let quit_button = Button::new("Quit", |s| s.quit());
/// ```
pub struct Button {
    label: String,
    callback: Callback,
    enabled: bool,
}

impl Button {
    /// Creates a new button with the given content and callback.
    pub fn new<F, R, S: Into<String>>(label: S, cb: F) -> Self
    where
        F: 'static + Fn(&mut Cursive) -> R,
    {
        let label = label.into();
        Self::new_raw(format!("<{}>", label), cb)
    }

    /// Creates a new button without angle brackets.
    pub fn new_raw<F, R, S: Into<String>>(label: S, cb: F) -> Self
    where
        F: 'static + Fn(&mut Cursive) -> R,
    {
        Button {
            label: label.into(),
            callback: Callback::from_fn(cb),
            enabled: true,
        }
    }

    /// Sets the function to be called when the button is pressed.
    ///
    /// Replaces the previous callback.
    pub fn set_callback<F>(&mut self, cb: F)
    where
        F: Fn(&mut Cursive) + 'static,
    {
        self.callback = Callback::from_fn(cb);
    }

    /// Disables this view.
    ///
    /// A disabled view cannot be selected.
    pub fn disable(&mut self) {
        self.enabled = false;
    }

    /// Disables this view.
    ///
    /// Chainable variant.
    pub fn disabled(self) -> Self {
        self.with(Self::disable)
    }

    /// Re-enables this view.
    pub fn enable(&mut self) {
        self.enabled = true;
    }

    /// Enable or disable this view.
    pub fn set_enabled(&mut self, enabled: bool) {
        self.enabled = enabled;
    }

    /// Returns `true` if this view is enabled.
    pub fn is_enabled(&self) -> bool {
        self.enabled
    }

    fn req_size(&self) -> Vec2 {
        Vec2::new(self.label.width(), 1)
    }
}

impl View for Button {
    fn draw(&self, printer: &Printer) {
        if printer.size.x == 0 {
            return;
        }

        let style = if !self.enabled {
            ColorStyle::Secondary
        } else if !printer.focused {
            ColorStyle::Primary
        } else {
            ColorStyle::Highlight
        };

        let offset =
            HAlign::Center.get_offset(self.label.len(), printer.size.x);

        printer.with_color(style, |printer| {
            printer.print((offset, 0), &self.label);
        });
    }

    fn required_size(&mut self, _: Vec2) -> Vec2 {
        // Meh. Fixed size we are.
        self.req_size()
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        // eprintln!("{:?}", event);
        // eprintln!("{:?}", self.req_size());
        match event {
            // 10 is the ascii code for '\n', that is the return key
            Event::Key(Key::Enter) => {
                EventResult::Consumed(Some(self.callback.clone()))
            }
            Event::Mouse {
                event: MouseEvent::Release(MouseButton::Left),
                position,
                offset,
            } if position.fits_in_rect(offset, self.req_size()) =>
            {
                EventResult::Consumed(Some(self.callback.clone()))
            }
            _ => EventResult::Ignored,
        }
    }

    fn take_focus(&mut self, _: Direction) -> bool {
        self.enabled
    }
}
<commit_msg>Button: add a way to get its label<commit_after>use {Cursive, Printer, With};
use align::HAlign;
use direction::Direction;
use event::*;
use theme::ColorStyle;
use unicode_width::UnicodeWidthStr;
use vec::Vec2;
use view::View;

/// Simple text label with a callback when <Enter> is pressed.
///
/// A button shows its content in a single line and has a fixed size.
///
/// # Examples
///
/// ```
/// # use cursive::views::Button;
/// let quit_button = Button::new("Quit", |s| s.quit());
/// ```
pub struct Button {
    label: String,
    callback: Callback,
    enabled: bool,
}

impl Button {
    /// Creates a new button with the given content and callback.
    pub fn new<F, R, S: Into<String>>(label: S, cb: F) -> Self
    where
        F: 'static + Fn(&mut Cursive) -> R,
    {
        let label = label.into();
        Self::new_raw(format!("<{}>", label), cb)
    }

    /// Creates a new button without angle brackets.
    pub fn new_raw<F, R, S: Into<String>>(label: S, cb: F) -> Self
    where
        F: 'static + Fn(&mut Cursive) -> R,
    {
        Button {
            label: label.into(),
            callback: Callback::from_fn(cb),
            enabled: true,
        }
    }

    /// Sets the function to be called when the button is pressed.
    ///
    /// Replaces the previous callback.
    pub fn set_callback<F>(&mut self, cb: F)
    where
        F: Fn(&mut Cursive) + 'static,
    {
        self.callback = Callback::from_fn(cb);
    }

    /// Disables this view.
    ///
    /// A disabled view cannot be selected.
    pub fn disable(&mut self) {
        self.enabled = false;
    }

    /// Disables this view.
    ///
    /// Chainable variant.
    pub fn disabled(self) -> Self {
        self.with(Self::disable)
    }

    /// Re-enables this view.
    pub fn enable(&mut self) {
        self.enabled = true;
    }

    /// Enable or disable this view.
    pub fn set_enabled(&mut self, enabled: bool) {
        self.enabled = enabled;
    }

    /// Returns `true` if this view is enabled.
    pub fn is_enabled(&self) -> bool {
        self.enabled
    }

    /// Returns the label for this button.
    ///
    /// Includes brackets.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use cursive::views::Button;
    /// let button = Button::new("Quit", |s| s.quit());
    /// assert_eq!(button.label(), "<Quit>");
    /// ```
    pub fn label(&self) -> &str {
        &self.label
    }

    fn req_size(&self) -> Vec2 {
        Vec2::new(self.label.width(), 1)
    }
}

impl View for Button {
    fn draw(&self, printer: &Printer) {
        if printer.size.x == 0 {
            return;
        }

        let style = if !self.enabled {
            ColorStyle::Secondary
        } else if !printer.focused {
            ColorStyle::Primary
        } else {
            ColorStyle::Highlight
        };

        let offset =
            HAlign::Center.get_offset(self.label.len(), printer.size.x);

        printer.with_color(style, |printer| {
            printer.print((offset, 0), &self.label);
        });
    }

    fn required_size(&mut self, _: Vec2) -> Vec2 {
        // Meh. Fixed size we are.
        self.req_size()
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        // eprintln!("{:?}", event);
        // eprintln!("{:?}", self.req_size());
        match event {
            // 10 is the ascii code for '\n', that is the return key
            Event::Key(Key::Enter) => {
                EventResult::Consumed(Some(self.callback.clone()))
            }
            Event::Mouse {
                event: MouseEvent::Release(MouseButton::Left),
                position,
                offset,
            } if position.fits_in_rect(offset, self.req_size()) =>
            {
                EventResult::Consumed(Some(self.callback.clone()))
            }
            _ => EventResult::Ignored,
        }
    }

    fn take_focus(&mut self, _: Direction) -> bool {
        self.enabled
    }
}
<|endoftext|>


<commit_before>use consts::HEADER_BYTES;
use trunking::fields::TalkGroup;
use util::slice_u16;
use voice::crypto;

pub struct VoiceHeaderFields([u8; HEADER_BYTES]);

impl VoiceHeaderFields {
    pub fn new(buf: [u8; HEADER_BYTES]) -> Self { VoiceHeaderFields(buf) }

    pub fn crypto_init(&self) -> &[u8] { &self.0[..9] }
    pub fn mfg(&self) -> u8 { self.0[9] }

    pub fn crypto_alg(&self) -> crypto::CryptoAlgorithm {
        crypto::CryptoAlgorithm::from_bits(self.0[10])
    }

    pub fn crypto_key(&self) -> u16 { slice_u16(&self.0[11..]) }

    pub fn talk_group(&self) -> TalkGroup {
        TalkGroup::from_bits(slice_u16(&self.0[13..]))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use voice::crypto::CryptoAlgorithm::*;
    use trunking::fields::TalkGroup;

    #[test]
    fn test_header() {
        let h = VoiceHeaderFields::new([
            1, 2, 3, 4, 5, 6, 7, 8, 9,
            0b00000000,
            0b10000000,
            0b00000000,
            0b00000000,
            0b11111111,
            0b11111111,
        ]);

        assert_eq!(h.crypto_init(), &[1,2,3,4,5,6,7,8,9]);
        assert_eq!(h.mfg(), 0);
        assert_eq!(h.crypto_alg(), Unencrypted);
        assert_eq!(h.crypto_key(), 0);
        assert_eq!(h.talk_group(), TalkGroup::Everbody);
    }
}
<commit_msg>Include CryptoAlgorithm instead of crypto<commit_after>use consts::HEADER_BYTES;
use trunking::fields::TalkGroup;
use util::slice_u16;
use voice::crypto::CryptoAlgorithm;

pub struct VoiceHeaderFields([u8; HEADER_BYTES]);

impl VoiceHeaderFields {
    pub fn new(buf: [u8; HEADER_BYTES]) -> Self { VoiceHeaderFields(buf) }

    pub fn crypto_init(&self) -> &[u8] { &self.0[..9] }
    pub fn mfg(&self) -> u8 { self.0[9] }

    pub fn crypto_alg(&self) -> CryptoAlgorithm {
        CryptoAlgorithm::from_bits(self.0[10])
    }

    pub fn crypto_key(&self) -> u16 { slice_u16(&self.0[11..]) }

    pub fn talk_group(&self) -> TalkGroup {
        TalkGroup::from_bits(slice_u16(&self.0[13..]))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use voice::crypto::CryptoAlgorithm::*;
    use trunking::fields::TalkGroup;

    #[test]
    fn test_header() {
        let h = VoiceHeaderFields::new([
            1, 2, 3, 4, 5, 6, 7, 8, 9,
            0b00000000,
            0b10000000,
            0b00000000,
            0b00000000,
            0b11111111,
            0b11111111,
        ]);

        assert_eq!(h.crypto_init(), &[1,2,3,4,5,6,7,8,9]);
        assert_eq!(h.mfg(), 0);
        assert_eq!(h.crypto_alg(), Unencrypted);
        assert_eq!(h.crypto_key(), 0);
        assert_eq!(h.talk_group(), TalkGroup::Everbody);
    }
}
<|endoftext|>


<commit_before>use std::collections::BinaryHeap;

use event::{self, Event};
use platform::Platform;
use profile::Profile;
use schedule::{self, Schedule};
use traffic::Traffic;
use workload::{Pattern, Workload};
use {Config, Result, Source};

pub struct System {
    platform: Platform,
    schedule: Box<Schedule>,
    traffic: Traffic,
    workload: Workload,
    stats: Stats,
    queue: BinaryHeap<Event>,
}

#[derive(Clone, Debug)]
pub struct Job {
    id: usize,
    arrival: f64,
    pattern: Pattern,
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Stats {
    pub created: usize,
    pub arrived: usize,
    pub started: usize,
    pub finished: usize,
}

pub type Increment = (Event, Profile, Profile);

impl System {
    pub fn new(config: Config, source: Source) -> Result<System> {
        let platform = {
            let config = some!(config.branch("platform"), "a platform configuration is required");
            try!(Platform::new(&config))
        };
        let schedule = {
            Box::new(try!(schedule::Compact::new(platform.elements())))
        };
        let traffic = {
            let config = some!(config.branch("traffic"), "a traffic configuration is required");
            try!(Traffic::new(&config, &source))
        };
        let workload = {
            let config = some!(config.branch("workload"), "a workload configuration is required");
            try!(Workload::new(&config, &source))
        };

        Ok(System {
            platform: platform,
            schedule: schedule,
            traffic: traffic,
            workload: workload,
            stats: Stats::default(),
            queue: BinaryHeap::new(),
        })
    }

    getters! {
        ref platform: Platform,
        ref stats: Stats,
    }
}

impl System {
    fn tick(&mut self) -> Result<()> {
        match (self.traffic.peek(), self.queue.peek()) {
            (Some(_), None) => {}
            (Some(&arrival), Some(&Event { time, .. })) if arrival < time => {},
            _ => return Ok(()),
        }

        let job = match (self.traffic.next(), self.workload.next()) {
            (Some(arrival), Some(pattern)) => {
                let id = self.stats.created;
                self.stats.created += 1;
                Job { id: id, arrival: arrival, pattern: pattern }
            },
            _ => raise!("failed to generate a new job"),
        };

        self.queue.push(Event::new(job.arrival, event::Kind::Arrival, job.clone()));

        let decision = try!(self.schedule.push(&job));
        try!(self.platform.push(&job, &decision));

        self.queue.push(Event::new(decision.start, event::Kind::Start, job.clone()));
        self.queue.push(Event::new(decision.finish, event::Kind::Finish, job));

        Ok(())
    }
}

impl Iterator for System {
    type Item = Increment;

    fn next(&mut self) -> Option<Self::Item> {
        if let Err(error) = self.tick() {
            error!(target: "System", "Failed to update the state ({}).", error);
        }
        let event = match self.queue.pop() {
            Some(event) => event,
            _ => return None,
        };
        self.schedule.pass(event.time);
        let (power, temperature) = match self.platform.next(event.time) {
            Some((power, temperature)) => (power, temperature),
            _ => return None,
        };
        self.stats.account(&event);
        Some((event, power, temperature))
    }
}

impl Job {
    getters! {
        id: usize,
        arrival: f64,
        ref pattern: Pattern,
    }
}

impl Stats {
    fn account(&mut self, event: &Event) {
        match event.kind {
            event::Kind::Arrival => self.arrived += 1,
            event::Kind::Start => self.started += 1,
            event::Kind::Finish => self.finished += 1,
        }
    }
}
<commit_msg>A cosmetic change<commit_after>use std::collections::BinaryHeap;

use event::Event;
use platform::Platform;
use profile::Profile;
use schedule::{self, Schedule};
use traffic::Traffic;
use workload::{Pattern, Workload};
use {Config, Result, Source};

pub struct System {
    platform: Platform,
    schedule: Box<Schedule>,
    traffic: Traffic,
    workload: Workload,
    stats: Stats,
    queue: BinaryHeap<Event>,
}

#[derive(Clone, Debug)]
pub struct Job {
    id: usize,
    arrival: f64,
    pattern: Pattern,
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Stats {
    pub created: usize,
    pub arrived: usize,
    pub started: usize,
    pub finished: usize,
}

pub type Increment = (Event, Profile, Profile);

impl System {
    pub fn new(config: Config, source: Source) -> Result<System> {
        let platform = {
            let config = some!(config.branch("platform"), "a platform configuration is required");
            try!(Platform::new(&config))
        };
        let schedule = {
            Box::new(try!(schedule::Compact::new(platform.elements())))
        };
        let traffic = {
            let config = some!(config.branch("traffic"), "a traffic configuration is required");
            try!(Traffic::new(&config, &source))
        };
        let workload = {
            let config = some!(config.branch("workload"), "a workload configuration is required");
            try!(Workload::new(&config, &source))
        };

        Ok(System {
            platform: platform,
            schedule: schedule,
            traffic: traffic,
            workload: workload,
            stats: Stats::default(),
            queue: BinaryHeap::new(),
        })
    }

    getters! {
        ref platform: Platform,
        ref stats: Stats,
    }
}

impl System {
    fn tick(&mut self) -> Result<()> {
        use event::Kind::*;

        match (self.traffic.peek(), self.queue.peek()) {
            (Some(_), None) => {}
            (Some(&arrival), Some(&Event { time, .. })) if arrival < time => {},
            _ => return Ok(()),
        }

        let job = match (self.traffic.next(), self.workload.next()) {
            (Some(arrival), Some(pattern)) => {
                let id = self.stats.created;
                self.stats.created += 1;
                Job { id: id, arrival: arrival, pattern: pattern }
            },
            _ => raise!("failed to generate a new job"),
        };

        self.queue.push(Event::new(job.arrival, Arrival, job.clone()));

        let decision = try!(self.schedule.push(&job));
        try!(self.platform.push(&job, &decision));

        self.queue.push(Event::new(decision.start, Start, job.clone()));
        self.queue.push(Event::new(decision.finish, Finish, job));

        Ok(())
    }
}

impl Iterator for System {
    type Item = Increment;

    fn next(&mut self) -> Option<Self::Item> {
        if let Err(error) = self.tick() {
            error!(target: "System", "Failed to update the state ({}).", error);
        }
        let event = match self.queue.pop() {
            Some(event) => event,
            _ => return None,
        };
        self.schedule.pass(event.time);
        let (power, temperature) = match self.platform.next(event.time) {
            Some((power, temperature)) => (power, temperature),
            _ => return None,
        };
        self.stats.account(&event);
        Some((event, power, temperature))
    }
}

impl Job {
    getters! {
        id: usize,
        arrival: f64,
        ref pattern: Pattern,
    }
}

impl Stats {
    fn account(&mut self, event: &Event) {
        use event::Kind::*;
        match event.kind {
            Arrival => self.arrived += 1,
            Start => self.started += 1,
            Finish => self.finished += 1,
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2016, 2020 É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/>.

use crate::error::{Error, Result};

use std::fs::{self, DirBuilder, File};
use std::io;
use std::io::Write;
use std::ops::Drop;
use std::path::{Path, PathBuf};
use std::process::Command;
use uuid;

/// Struct used to create zip (using filesystem and zip command)
pub struct Zipper {
    args: Vec<String>,
    path: PathBuf,
}

impl Zipper {
    /// Creates new zipper
    ///
    /// # Arguments
    /// * `path`: the path to a temporary directory
    /// (zipper will create a random dir in it and clean it later)
    pub fn new(path: &str) -> Result<Zipper> {
        let uuid = uuid::Uuid::new_v4();
        let zipper_path = Path::new(path).join(uuid.to_simple().to_string());

        DirBuilder::new()
            .recursive(true)
            .create(&zipper_path)
            .map_err(|_| {
                Error::zipper(lformat!(
                    "could not create temporary directory in {path}",
                    path = path
                ))
            })?;

        Ok(Zipper {
            args: vec![],
            path: zipper_path,
        })
    }

    /// writes a content to a temporary file
    pub fn write<P: AsRef<Path>>(&mut self, path: P, content: &[u8], add_args: bool) -> Result<()> {
        let path = path.as_ref();
        let file = format!("{}", path.display());
        if path.starts_with("..") || path.is_absolute() {
            return Err(Error::zipper(lformat!(
                "file {file} refers to an absolute or a parent \
                                               path.
This is forbidden because we are supposed \
                                               to create a temporary file in a temporary dir.",
                file = file
            )));
        }
        let dest_file = self.path.join(path);
        let dest_dir = dest_file.parent().unwrap();
        if !fs::metadata(dest_dir).is_ok() {
            // dir does not exist, create it
            DirBuilder::new()
                .recursive(true)
                .create(&dest_dir)
                .map_err(|_| {
                    Error::zipper(lformat!(
                        "could not create temporary directory in {path}",
                        path = dest_dir.display()
                    ))
                })?;
        }

        if let Ok(mut f) = File::create(&dest_file) {
            if f.write_all(content).is_ok() {
                if add_args {
                    self.args.push(String::from(file));
                }
                Ok(())
            } else {
                Err(Error::zipper(lformat!(
                    "could not write to temporary file {file}",
                    file = file
                )))
            }
        } else {
            Err(Error::zipper(lformat!(
                "could not create temporary file {file}",
                file = file
            )))
        }
    }

    /// Unzip a file and deletes it afterwards
    pub fn unzip(&mut self, file: &str) -> Result<()> {
        let output = Command::new("unzip")
            .current_dir(&self.path)
            .arg(file)
            .output()
            .map_err(|e| {
                Error::zipper(lformat!(
                    "failed to execute unzip on {file}: {error}",
                    file = file,
                    error = e
                ))
            });

        output?;

        fs::remove_file(self.path.join(file))
            .map_err(|_| Error::zipper(lformat!("failed to remove file {file}", file = file)))
    }

    /// run command and copy content of file output (supposed to result from the command) to current dir
    pub fn run_command(
        &mut self,
        mut command: Command,
        command_name: &str,
        in_file: &str,
        out: &mut dyn Write,
    ) -> Result<String> {
        let res_output = command
            .current_dir(&self.path)
            .output()
            .map_err(|e| {
                debug!(
                    "{}",
                    lformat!(
                        "output for command {name}:\n{error}",
                        name = command_name,
                        error = e
                    )
                );
                Error::zipper(lformat!(
                    "failed to run command '{name}'",
                    name = command_name
                ))
            });
        let output = res_output?;
        if output.status.success() {
            let mut file = File::open(self.path.join(in_file)).map_err(|_| {
                debug!(
                    "{}",
                    lformat!(
                        "could not open result of command '{command}'\n\
                                           Command output:\n\
                                           {output}'",
                        command = command_name,
                        output = String::from_utf8_lossy(&output.stdout)
                    )
                );
                Error::zipper(lformat!(
                    "could not open result of command '{command}'",
                    command = command_name
                ))
            })?;
            io::copy(&mut file, out).map_err(|_| {
                Error::zipper(lformat!("error copying file '{file}'", file = in_file))
            })?;

            Ok(String::from_utf8_lossy(&output.stdout).into_owned())
        } else {
            debug!(
                "{}",
                lformat!(
                    "{command} didn't return succesfully: {output}",
                    command = command_name,
                    output = String::from_utf8_lossy(&output.stdout)
                )
            );
            Err(Error::zipper(lformat!(
                "{command} didn't return succesfully",
                command = command_name
            )))
        }
    }

    /// zip all files in zipper's tmp dir to a given file name and write to odt file
    pub fn generate_odt(&mut self, command_name: &str, odt_file: &mut dyn Write) -> Result<String> {
        let mut command = Command::new(command_name);
        command.arg("-r");
        command.arg("result.odt");
        command.arg(".");
        self.run_command(command, command_name, "result.odt", odt_file)
    }

    /// generate a pdf file into given file name
    pub fn generate_pdf(
        &mut self,
        command_name: &str,
        tex_file: &str,
        pdf_file: &mut dyn Write,
    ) -> Result<String> {
        // first pass
        let mut command = Command::new(command_name);
        command.current_dir(&self.path).arg(tex_file);
        let _ = command.output();

        // second pass
        let _ = command.output();

        // third pass
        // let mut command = Command::new(command_name);
        // command.current_dir(&self.path);
        // command.arg(tex_file);
        self.run_command(command, command_name, "result.pdf", pdf_file)
    }
}

impl Drop for Zipper {
    fn drop(&mut self) {
        if let Err(err) = fs::remove_dir_all(&self.path) {
            println!(
                "Error in zipper: could not delete temporary directory {}, error: {}",
                self.path.to_string_lossy(),
                err
            );
        }
    }
}
<commit_msg>run_command does not need to set the path again<commit_after>// Copyright (C) 2016, 2020 É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/>.

use crate::error::{Error, Result};

use std::fs::{self, DirBuilder, File};
use std::io;
use std::io::Write;
use std::ops::Drop;
use std::path::{Path, PathBuf};
use std::process::Command;
use uuid;

/// Struct used to create zip (using filesystem and zip command)
pub struct Zipper {
    args: Vec<String>,
    path: PathBuf,
}

impl Zipper {
    /// Creates new zipper
    ///
    /// # Arguments
    /// * `path`: the path to a temporary directory
    /// (zipper will create a random dir in it and clean it later)
    pub fn new(path: &str) -> Result<Zipper> {
        let uuid = uuid::Uuid::new_v4();
        let zipper_path = Path::new(path).join(uuid.to_simple().to_string());

        DirBuilder::new()
            .recursive(true)
            .create(&zipper_path)
            .map_err(|_| {
                Error::zipper(lformat!(
                    "could not create temporary directory in {path}",
                    path = path
                ))
            })?;

        Ok(Zipper {
            args: vec![],
            path: zipper_path,
        })
    }

    /// writes a content to a temporary file
    pub fn write<P: AsRef<Path>>(&mut self, path: P, content: &[u8], add_args: bool) -> Result<()> {
        let path = path.as_ref();
        let file = format!("{}", path.display());
        if path.starts_with("..") || path.is_absolute() {
            return Err(Error::zipper(lformat!(
                "file {file} refers to an absolute or a parent \
                                               path.
This is forbidden because we are supposed \
                                               to create a temporary file in a temporary dir.",
                file = file
            )));
        }
        let dest_file = self.path.join(path);
        let dest_dir = dest_file.parent().unwrap();
        if !fs::metadata(dest_dir).is_ok() {
            // dir does not exist, create it
            DirBuilder::new()
                .recursive(true)
                .create(&dest_dir)
                .map_err(|_| {
                    Error::zipper(lformat!(
                        "could not create temporary directory in {path}",
                        path = dest_dir.display()
                    ))
                })?;
        }

        if let Ok(mut f) = File::create(&dest_file) {
            if f.write_all(content).is_ok() {
                if add_args {
                    self.args.push(String::from(file));
                }
                Ok(())
            } else {
                Err(Error::zipper(lformat!(
                    "could not write to temporary file {file}",
                    file = file
                )))
            }
        } else {
            Err(Error::zipper(lformat!(
                "could not create temporary file {file}",
                file = file
            )))
        }
    }

    /// Unzip a file and deletes it afterwards
    pub fn unzip(&mut self, file: &str) -> Result<()> {
        let output = Command::new("unzip")
            .current_dir(&self.path)
            .arg(file)
            .output()
            .map_err(|e| {
                Error::zipper(lformat!(
                    "failed to execute unzip on {file}: {error}",
                    file = file,
                    error = e
                ))
            });

        output?;

        fs::remove_file(self.path.join(file))
            .map_err(|_| Error::zipper(lformat!("failed to remove file {file}", file = file)))
    }

    /// run command and copy content of file output (supposed to result from the command) to current dir
    pub fn run_command(
        &mut self,
        mut command: Command,
        command_name: &str,
        in_file: &str,
        out: &mut dyn Write,
    ) -> Result<String> {
        let res_output = command.output().map_err(|e| {
            debug!(
                "{}",
                lformat!(
                    "output for command {name}:\n{error}",
                    name = command_name,
                    error = e
                )
            );
            Error::zipper(lformat!(
                "failed to run command '{name}'",
                name = command_name
            ))
        });
        let output = res_output?;
        if output.status.success() {
            let mut file = File::open(self.path.join(in_file)).map_err(|_| {
                debug!(
                    "{}",
                    lformat!(
                        "could not open result of command '{command}'\n\
                                           Command output:\n\
                                           {output}'",
                        command = command_name,
                        output = String::from_utf8_lossy(&output.stdout)
                    )
                );
                Error::zipper(lformat!(
                    "could not open result of command '{command}'",
                    command = command_name
                ))
            })?;
            io::copy(&mut file, out).map_err(|_| {
                Error::zipper(lformat!("error copying file '{file}'", file = in_file))
            })?;

            Ok(String::from_utf8_lossy(&output.stdout).into_owned())
        } else {
            debug!(
                "{}",
                lformat!(
                    "{command} didn't return succesfully: {output}",
                    command = command_name,
                    output = String::from_utf8_lossy(&output.stdout)
                )
            );
            Err(Error::zipper(lformat!(
                "{command} didn't return succesfully",
                command = command_name
            )))
        }
    }

    /// zip all files in zipper's tmp dir to a given file name and write to odt file
    pub fn generate_odt(&mut self, command_name: &str, odt_file: &mut dyn Write) -> Result<String> {
        let mut command = Command::new(command_name);
        command.current_dir(&self.path);
        command.arg("-r");
        command.arg("result.odt");
        command.arg(".");
        self.run_command(command, command_name, "result.odt", odt_file)
    }

    /// generate a pdf file into given file name
    pub fn generate_pdf(
        &mut self,
        command_name: &str,
        tex_file: &str,
        pdf_file: &mut dyn Write,
    ) -> Result<String> {
        // first pass
        let mut command = Command::new(command_name);
        command.current_dir(&self.path).arg(tex_file);
        let _ = command.output();

        // second pass
        let _ = command.output();

        // third pass
        // let mut command = Command::new(command_name);
        // command.current_dir(&self.path);
        // command.arg(tex_file);
        self.run_command(command, command_name, "result.pdf", pdf_file)
    }
}

impl Drop for Zipper {
    fn drop(&mut self) {
        if let Err(err) = fs::remove_dir_all(&self.path) {
            println!(
                "Error in zipper: could not delete temporary directory {}, error: {}",
                self.path.to_string_lossy(),
                err
            );
        }
    }
}
<|endoftext|>


<commit_before>use scroll::{self, Pread};
use error::{self, Error};

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        for i in 0..8 {
            name[i] = bytes.gread_with(offset, scroll::LE)?;
        }

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if name[0] == b'/' {
            let idx: usize = if name[1] == b'/' {
                // TODO: Base-64 encoding
                panic!("At the disco")
            } else {
                let name = name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x00000008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x00000020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x00000040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x00000080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x00000100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x00000200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x00000800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x00001000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x00008000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x00040000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x00080000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x00100000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x00200000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x00300000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x00400000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x00500000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x00600000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x00700000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x00800000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x00900000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A00000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B00000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C00000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D00000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E00000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F00000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x01000000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x02000000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x04000000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x08000000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x10000000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x40000000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x80000000;
<commit_msg>Implement base64 section names<commit_after>use scroll::{self, Pread};
use error::{self, Error};

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

// Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70
// Decodes a string table entry in base 64 (//AAAAAA). Expects string without
// prefixed slashes.
fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
    assert!(s.len() <= 6, "String too long, possible overflow.");

    let mut val = 0;
    for c in s.bytes() {
        let v = match c {
            b'A'..=b'Z' => c - b'A' + 00, // 00..=25
            b'a'..=b'z' => c - b'a' + 26, // 26..=51
            b'0'..=b'9' => c - b'0' + 52, // 52..=61
            b'+'        => 62,            // 62
            b'/'        => 63,            // 63
            _           => return Err(())
        };
        val = val * 64 + v as usize;
    }
    Ok(val)
}

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        for i in 0..8 {
            name[i] = bytes.gread_with(offset, scroll::LE)?;
        }

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if name[0] == b'/' {
            let idx: usize = if name[1] == b'/' {
                let b64idx = name.pread::<&str>(2)?;
                base64_decode_string_entry(b64idx).map_err(|_|
                    Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
            } else {
                let name = name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x00000008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x00000020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x00000040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x00000080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x00000100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x00000200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x00000800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x00001000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x00008000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x00020000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x00040000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x00080000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x00100000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x00200000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x00300000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x00400000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x00500000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x00600000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x00700000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x00800000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x00900000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A00000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B00000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C00000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D00000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E00000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F00000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x01000000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x02000000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x04000000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x08000000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x10000000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x40000000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x80000000;
<|endoftext|>