<commit_before>use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::io::Read;

use crate::common::LineIter;

type Point = (i32, i32);

struct Map {
    width: usize,
    data: Vec<u8>,
}

impl Map {
    pub fn from_input(input: &mut dyn Read) -> Self {
        let mut reader = LineIter::new(input);

        let mut data = reader.next().unwrap().as_bytes().to_owned();
        let width = data.len();

        while let Some(line) = reader.next() {
            let line = line.as_bytes();
            debug_assert_eq!(line.len(), width);

            data.extend_from_slice(line);
        }

        data.iter_mut().for_each(|b| *b -= b'0');

        Self { width, data }
    }

    pub fn from_input2(input: &mut dyn Read) -> Self {
        let mut reader = LineIter::new(input);

        let mut lines = Vec::new();

        while let Some(line) = reader.next() {
            let mut line = line.as_bytes().to_owned();
            line.iter_mut().for_each(|b| *b -= b'0');
            lines.push(line);
        }

        let mut data = Vec::new();
        let width = lines[0].len();

        for _y_repeat in 0..5 {
            for line in &mut lines {
                data.extend_from_slice(line);

                for _ in 0..4 {
                    let starting_pos = data.len() - width;
                    data.extend_from_within(starting_pos..);
                    let starting_pos = data.len() - width;
                    data[starting_pos..]
                        .iter_mut()
                        .for_each(|b| *b = (*b % 9) + 1);
                }

                line.iter_mut().for_each(|b| *b = (*b % 9) + 1);
            }
        }

        Self {
            width: width * 5,
            data,
        }
    }

    pub fn shortest_path(&self, start: Point, end: Point) -> u32 {
        let mut todo = BinaryHeap::new();
        todo.push(Reverse((0, start)));

        let mut visited = vec![false; self.data.len()];

        let height = self.height() as i32;

        while let Some(Reverse((distance, pos))) = todo.pop() {
            if pos == end {
                return distance;
            }

            if visited[self.index(pos)] {
                continue;
            }

            visited[self.index(pos)] = true;

            let (x, y) = pos;

            for dy in -1..=1 {
                if y + dy < 0 || y + dy >= height {
                    continue;
                }

                for dx in -1..=1 {
                    if x + dx < 0 || (x + dx) >= self.width as i32 || dx * dy != 0 {
                        continue;
                    }

                    let new = (x + dx, y + dy);
                    let index = self.index(new);

                    if visited[index] {
                        continue;
                    }

                    let new_distance = distance + self.data[index] as u32;

                    todo.push(Reverse((new_distance, new)));
                }
            }
        }

        panic!("No route found from {:?} to {:?}", start, end);
    }

    pub fn height(&self) -> usize {
        self.data.len() / self.width
    }

    pub fn index(&self, (x, y): Point) -> usize {
        y as usize * self.width + x as usize
    }
}

pub fn part1(input: &mut dyn Read) -> String {
    let map = Map::from_input(input);

    map.shortest_path((0, 0), (map.width as i32 - 1, map.height() as i32 - 1))
        .to_string()
}

pub fn part2(input: &mut dyn Read) -> String {
    let map = Map::from_input2(input);

    map.shortest_path((0, 0), (map.width as i32 - 1, map.height() as i32 - 1))
        .to_string()
}

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

    use crate::test_implementation;

    const SAMPLE: &[u8] = include_bytes!("samples/15.txt");

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 40);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 315);
    }
}
<commit_msg>Avoid queueing worse routes<commit_after>use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::io::Read;

use crate::common::LineIter;

type Point = (i32, i32);

struct Map {
    width: usize,
    data: Vec<u8>,
}

impl Map {
    pub fn from_input(input: &mut dyn Read) -> Self {
        let mut reader = LineIter::new(input);

        let mut data = reader.next().unwrap().as_bytes().to_owned();
        let width = data.len();

        while let Some(line) = reader.next() {
            let line = line.as_bytes();
            debug_assert_eq!(line.len(), width);

            data.extend_from_slice(line);
        }

        data.iter_mut().for_each(|b| *b -= b'0');

        Self { width, data }
    }

    pub fn from_input2(input: &mut dyn Read) -> Self {
        let mut reader = LineIter::new(input);

        let mut lines = Vec::new();

        while let Some(line) = reader.next() {
            let mut line = line.as_bytes().to_owned();
            line.iter_mut().for_each(|b| *b -= b'0');
            lines.push(line);
        }

        let mut data = Vec::new();
        let width = lines[0].len();

        for _y_repeat in 0..5 {
            for line in &mut lines {
                data.extend_from_slice(line);

                for _ in 0..4 {
                    let starting_pos = data.len() - width;
                    data.extend_from_within(starting_pos..);
                    let starting_pos = data.len() - width;
                    data[starting_pos..]
                        .iter_mut()
                        .for_each(|b| *b = (*b % 9) + 1);
                }

                line.iter_mut().for_each(|b| *b = (*b % 9) + 1);
            }
        }

        Self {
            width: width * 5,
            data,
        }
    }

    pub fn shortest_path(&self, start: Point, end: Point) -> u32 {
        let mut todo = BinaryHeap::new();
        todo.push(Reverse((0, start)));

        let mut best = vec![u32::MAX; self.data.len()];

        let height = self.height() as i32;

        while let Some(Reverse((distance, pos))) = todo.pop() {
            if pos == end {
                return distance;
            }

            if best[self.index(pos)] < distance {
                continue;
            }

            let (x, y) = pos;

            for dy in -1..=1 {
                if y + dy < 0 || y + dy >= height {
                    continue;
                }

                for dx in -1..=1 {
                    if x + dx < 0 || (x + dx) >= self.width as i32 || dx * dy != 0 {
                        continue;
                    }

                    let new = (x + dx, y + dy);
                    let index = self.index(new);
                    let new_distance = distance + self.data[index] as u32;

                    if best[index] <= new_distance {
                        continue;
                    }

                    best[index] = new_distance;

                    todo.push(Reverse((new_distance, new)));
                }
            }
        }

        panic!("No route found from {:?} to {:?}", start, end);
    }

    pub fn height(&self) -> usize {
        self.data.len() / self.width
    }

    pub fn index(&self, (x, y): Point) -> usize {
        y as usize * self.width + x as usize
    }
}

pub fn part1(input: &mut dyn Read) -> String {
    let map = Map::from_input(input);

    map.shortest_path((0, 0), (map.width as i32 - 1, map.height() as i32 - 1))
        .to_string()
}

pub fn part2(input: &mut dyn Read) -> String {
    let map = Map::from_input2(input);

    map.shortest_path((0, 0), (map.width as i32 - 1, map.height() as i32 - 1))
        .to_string()
}

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

    use crate::test_implementation;

    const SAMPLE: &[u8] = include_bytes!("samples/15.txt");

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 40);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 315);
    }
}
<|endoftext|>


<commit_before>use roaring::{RoaringBitmap, Iter};

use std::collections::{BTreeMap, HashMap};
use std::collections::btree_map::Keys;

use document::Document;
use schema::{Schema, SchemaRead, SchemaWrite, FieldType, FieldRef, AddFieldError};
use store::{IndexStore, IndexReader, DocRefIterator};

#[derive(Debug)]
pub struct MemoryIndexStoreFieldTerm {
    pub docs: RoaringBitmap<u64>,
}

impl MemoryIndexStoreFieldTerm {
    pub fn new() -> MemoryIndexStoreFieldTerm {
        MemoryIndexStoreFieldTerm {
            docs: RoaringBitmap::new(),
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStoreField {
    pub docs: RoaringBitmap<u64>,
    pub terms: BTreeMap<Vec<u8>, MemoryIndexStoreFieldTerm>,
    pub num_tokens: u64,
}

impl MemoryIndexStoreField {
    pub fn new() -> MemoryIndexStoreField {
        MemoryIndexStoreField {
            docs: RoaringBitmap::new(),
            terms: BTreeMap::new(),
            num_tokens: 0,
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStore {
    schema: Schema,
    docs: BTreeMap<u64, Document>,
    fields: HashMap<FieldRef, MemoryIndexStoreField>,
    next_doc_id: u64,
    doc_key2id_map: HashMap<String, u64>,
}

impl MemoryIndexStore {
    pub fn new() -> MemoryIndexStore {
        MemoryIndexStore {
            schema: Schema::new(),
            docs: BTreeMap::new(),
            fields: HashMap::new(),
            next_doc_id: 1,
            doc_key2id_map: HashMap::new(),
        }
    }
}

impl<'a> IndexStore<'a> for MemoryIndexStore {
    type Reader = MemoryIndexStoreReader<'a>;

    fn reader(&'a self) -> MemoryIndexStoreReader<'a> {
        MemoryIndexStoreReader {
            store: self,
        }
    }

    fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let field_ref = try!(self.schema.add_field(name, field_type));
        self.fields.insert(field_ref, MemoryIndexStoreField::new());

        Ok(field_ref)
    }

    fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        self.schema.remove_field(field_ref);

        match self.fields.remove(field_ref) {
            Some(field_ref) => true,
            None => false,
        }
    }

    fn insert_or_update_document(&mut self, doc: Document) {
        let doc_id = self.next_doc_id;
        self.next_doc_id += 1;

        // Put field contents in inverted index
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                // Silently ignore unrecognised fields
                // TODO: Review this
                if let Some(field) = self.fields.get_mut(field_ref) {
                    field.docs.insert(doc_id);
                    field.num_tokens += 1;

                    let term_bytes = token.term.to_bytes();
                    if !field.terms.contains_key(&term_bytes) {
                        // TODO: We shouldn't need to clone here
                        field.terms.insert(term_bytes.clone(), MemoryIndexStoreFieldTerm::new());
                    }

                    let mut term = field.terms.get_mut(&term_bytes).unwrap();
                    term.docs.insert(doc_id);
                }
            }
        }

        self.doc_key2id_map.insert(doc.key.clone(), doc_id);
        self.docs.insert(doc_id, doc);
    }

    fn remove_document_by_key(&mut self, doc_key: &str) -> bool {
        match self.doc_key2id_map.remove(doc_key) {
            Some(doc_id) => {
                self.docs.remove(&doc_id);

                true
            }
            None => false
        }
    }
}

pub struct MemoryIndexStoreReader<'a> {
    store: &'a MemoryIndexStore,
}

impl<'a> IndexReader<'a> for MemoryIndexStoreReader<'a> {
    type AllDocRefIterator = MemoryIndexStoreAllDocRefIterator<'a>;
    type TermDocRefIterator = MemoryIndexStoreTermDocRefIterator<'a>;

    fn schema(&self) -> &Schema {
        &self.store.schema
    }

    fn get_document_by_key(&self, doc_key: &str) -> Option<&Document> {
        match self.store.doc_key2id_map.get(doc_key) {
            Some(doc_id) => self.store.docs.get(doc_id),
            None => None,
        }
    }

    fn get_document_by_id(&self, doc_id: &u64) -> Option<&Document> {
        self.store.docs.get(doc_id)
    }

    fn contains_document_key(&self, doc_key: &str) -> bool {
        self.store.doc_key2id_map.contains_key(doc_key)
    }

    fn num_docs(&self) -> usize {
        self.store.docs.len()
    }

    fn iter_all_docs(&'a self) -> MemoryIndexStoreAllDocRefIterator<'a> {
        MemoryIndexStoreAllDocRefIterator {
            keys: self.store.docs.keys(),
        }
    }

    fn iter_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> Option<MemoryIndexStoreTermDocRefIterator<'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return None,
        };

        Some(MemoryIndexStoreTermDocRefIterator {
            iterator: term.docs.iter(),
        })
    }

    fn iter_all_terms(&'a self, field_ref: &FieldRef) -> Option<Box<Iterator<Item=&'a [u8]> + 'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        Some(Box::new(field.terms.keys().map(|t| &t[..])))
    }

    fn num_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return 0,
        };

        term.docs.len()
    }

    fn total_tokens(&'a self, field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        field.num_tokens
    }
}

pub struct MemoryIndexStoreAllDocRefIterator<'a> {
    keys: Keys<'a, u64, Document>,
}

impl<'a> Iterator for MemoryIndexStoreAllDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.keys.next().cloned()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreAllDocRefIterator<'a> {

}

pub struct MemoryIndexStoreTermDocRefIterator<'a> {
    iterator: Iter<'a, u64>,
}

impl<'a> Iterator for MemoryIndexStoreTermDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.iterator.next()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreTermDocRefIterator<'a> {

}

#[cfg(test)]
mod tests {
    use super::{MemoryIndexStore, MemoryIndexStoreReader};

    use term::Term;
    use token::Token;
    use document::Document;
    use schema::{Schema, SchemaRead, FieldType, FieldRef};
    use store::{IndexStore, IndexReader};

    fn make_test_store() -> MemoryIndexStore {
        let mut store = MemoryIndexStore::new();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    #[test]
    fn test_num_docs() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.num_docs(), 2);
    }

    #[test]
    fn test_all_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.iter_all_docs().count(), 2);
    }

    #[test]
    fn test_term_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();
        let title_field = reader.schema().get_field_by_name("title").unwrap();

        assert_eq!(reader.iter_docs_with_term(&Term::String("hello".to_string()).to_bytes(), &title_field).unwrap().count(), 1);
    }
}
<commit_msg>Improve performance of insert_or_update_document<commit_after>use roaring::{RoaringBitmap, Iter};

use std::collections::{BTreeMap, HashMap};
use std::collections::btree_map::Keys;

use document::Document;
use schema::{Schema, SchemaRead, SchemaWrite, FieldType, FieldRef, AddFieldError};
use store::{IndexStore, IndexReader, DocRefIterator};

#[derive(Debug)]
pub struct MemoryIndexStoreFieldTerm {
    pub docs: RoaringBitmap<u64>,
}

impl MemoryIndexStoreFieldTerm {
    pub fn new() -> MemoryIndexStoreFieldTerm {
        MemoryIndexStoreFieldTerm {
            docs: RoaringBitmap::new(),
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStoreField {
    pub docs: RoaringBitmap<u64>,
    pub terms: BTreeMap<Vec<u8>, MemoryIndexStoreFieldTerm>,
    pub num_tokens: u64,
}

impl MemoryIndexStoreField {
    pub fn new() -> MemoryIndexStoreField {
        MemoryIndexStoreField {
            docs: RoaringBitmap::new(),
            terms: BTreeMap::new(),
            num_tokens: 0,
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStore {
    schema: Schema,
    docs: BTreeMap<u64, Document>,
    fields: HashMap<FieldRef, MemoryIndexStoreField>,
    next_doc_id: u64,
    doc_key2id_map: HashMap<String, u64>,
}

impl MemoryIndexStore {
    pub fn new() -> MemoryIndexStore {
        MemoryIndexStore {
            schema: Schema::new(),
            docs: BTreeMap::new(),
            fields: HashMap::new(),
            next_doc_id: 1,
            doc_key2id_map: HashMap::new(),
        }
    }
}

impl<'a> IndexStore<'a> for MemoryIndexStore {
    type Reader = MemoryIndexStoreReader<'a>;

    fn reader(&'a self) -> MemoryIndexStoreReader<'a> {
        MemoryIndexStoreReader {
            store: self,
        }
    }

    fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let field_ref = try!(self.schema.add_field(name, field_type));
        self.fields.insert(field_ref, MemoryIndexStoreField::new());

        Ok(field_ref)
    }

    fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        self.schema.remove_field(field_ref);

        match self.fields.remove(field_ref) {
            Some(field_ref) => true,
            None => false,
        }
    }

    fn insert_or_update_document(&mut self, doc: Document) {
        let doc_id = self.next_doc_id;
        self.next_doc_id += 1;

        // Put field contents in inverted index
        for (field_ref, tokens) in doc.fields.iter() {
            // Silently ignore unrecognised fields
            // TODO: Review this
            if let Some(field) = self.fields.get_mut(field_ref) {
                field.docs.insert(doc_id);

                for token in tokens.iter() {
                    field.num_tokens += 1;

                    let term_bytes = token.term.to_bytes();
                    let mut term = field.terms.entry(term_bytes).or_insert_with(|| MemoryIndexStoreFieldTerm::new());
                    term.docs.insert(doc_id);
                }
            }
        }

        self.doc_key2id_map.insert(doc.key.clone(), doc_id);
        self.docs.insert(doc_id, doc);
    }

    fn remove_document_by_key(&mut self, doc_key: &str) -> bool {
        match self.doc_key2id_map.remove(doc_key) {
            Some(doc_id) => {
                self.docs.remove(&doc_id);

                true
            }
            None => false
        }
    }
}

pub struct MemoryIndexStoreReader<'a> {
    store: &'a MemoryIndexStore,
}

impl<'a> IndexReader<'a> for MemoryIndexStoreReader<'a> {
    type AllDocRefIterator = MemoryIndexStoreAllDocRefIterator<'a>;
    type TermDocRefIterator = MemoryIndexStoreTermDocRefIterator<'a>;

    fn schema(&self) -> &Schema {
        &self.store.schema
    }

    fn get_document_by_key(&self, doc_key: &str) -> Option<&Document> {
        match self.store.doc_key2id_map.get(doc_key) {
            Some(doc_id) => self.store.docs.get(doc_id),
            None => None,
        }
    }

    fn get_document_by_id(&self, doc_id: &u64) -> Option<&Document> {
        self.store.docs.get(doc_id)
    }

    fn contains_document_key(&self, doc_key: &str) -> bool {
        self.store.doc_key2id_map.contains_key(doc_key)
    }

    fn num_docs(&self) -> usize {
        self.store.docs.len()
    }

    fn iter_all_docs(&'a self) -> MemoryIndexStoreAllDocRefIterator<'a> {
        MemoryIndexStoreAllDocRefIterator {
            keys: self.store.docs.keys(),
        }
    }

    fn iter_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> Option<MemoryIndexStoreTermDocRefIterator<'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return None,
        };

        Some(MemoryIndexStoreTermDocRefIterator {
            iterator: term.docs.iter(),
        })
    }

    fn iter_all_terms(&'a self, field_ref: &FieldRef) -> Option<Box<Iterator<Item=&'a [u8]> + 'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        Some(Box::new(field.terms.keys().map(|t| &t[..])))
    }

    fn num_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return 0,
        };

        term.docs.len()
    }

    fn total_tokens(&'a self, field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        field.num_tokens
    }
}

pub struct MemoryIndexStoreAllDocRefIterator<'a> {
    keys: Keys<'a, u64, Document>,
}

impl<'a> Iterator for MemoryIndexStoreAllDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.keys.next().cloned()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreAllDocRefIterator<'a> {

}

pub struct MemoryIndexStoreTermDocRefIterator<'a> {
    iterator: Iter<'a, u64>,
}

impl<'a> Iterator for MemoryIndexStoreTermDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.iterator.next()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreTermDocRefIterator<'a> {

}

#[cfg(test)]
mod tests {
    use super::{MemoryIndexStore, MemoryIndexStoreReader};

    use term::Term;
    use token::Token;
    use document::Document;
    use schema::{Schema, SchemaRead, FieldType, FieldRef};
    use store::{IndexStore, IndexReader};

    fn make_test_store() -> MemoryIndexStore {
        let mut store = MemoryIndexStore::new();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    #[test]
    fn test_num_docs() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.num_docs(), 2);
    }

    #[test]
    fn test_all_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.iter_all_docs().count(), 2);
    }

    #[test]
    fn test_term_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();
        let title_field = reader.schema().get_field_by_name("title").unwrap();

        assert_eq!(reader.iter_docs_with_term(&Term::String("hello".to_string()).to_bytes(), &title_field).unwrap().count(), 1);
    }
}
<|endoftext|>


<commit_before>//! This file contains the memory allocator used by the kernel. It is a thin wrapper around
//! smallheap.

use alloc::heap::{Alloc, AllocErr, Layout};
use core::cell::RefCell;

use smallheap::Allocator;

use interrupts::{off, on};

/// A wrapper around the heap allocator for use as the `global_allocator`.
pub struct KernelAllocator {
    heap: RefCell<Option<Allocator>>,
}

impl KernelAllocator {
    pub const fn new() -> Self {
        KernelAllocator { heap: RefCell::new(None) }
    }

    pub fn set_heap(&mut self, heap: Allocator) {
        *self.heap.borrow_mut() = Some(heap);
    }
}

unsafe impl<'a> Alloc for &'a KernelAllocator {
    unsafe fn alloc(&mut self, layout: Layout) -> Result<*mut u8, AllocErr> {
        off();
        let ret = self.heap.borrow_mut().as_mut().unwrap().malloc(layout.size(), layout.align());
        on();
        Ok(ret)
    }

    unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
        off();
        let ret = self.heap.borrow_mut().as_mut().unwrap().free(ptr, layout.size());
        on();
        ret
    }
}

/// Initialize the kernel heap
pub fn init(allocator: &mut KernelAllocator, start: usize, size: usize) {
    let heap = unsafe {Allocator::new(start as *mut u8, size)};
    let free_size = heap.size();

    allocator.set_heap(heap);

    bootlog!(
        "heap inited - start addr: 0x{:x}, end addr: 0x{:x}, {} bytes\n",
        start,
        start + size,
        free_size
    );
}
<commit_msg>rustfmt<commit_after>//! This file contains the memory allocator used by the kernel. It is a thin wrapper around
//! smallheap.

use alloc::heap::{Alloc, AllocErr, Layout};
use core::cell::RefCell;

use smallheap::Allocator;

use interrupts::{off, on};

/// A wrapper around the heap allocator for use as the `global_allocator`.
pub struct KernelAllocator {
    heap: RefCell<Option<Allocator>>,
}

impl KernelAllocator {
    pub const fn new() -> Self {
        KernelAllocator {
            heap: RefCell::new(None),
        }
    }

    pub fn set_heap(&mut self, heap: Allocator) {
        *self.heap.borrow_mut() = Some(heap);
    }
}

unsafe impl<'a> Alloc for &'a KernelAllocator {
    unsafe fn alloc(&mut self, layout: Layout) -> Result<*mut u8, AllocErr> {
        off();
        let ret = self.heap
            .borrow_mut()
            .as_mut()
            .unwrap()
            .malloc(layout.size(), layout.align());
        on();
        Ok(ret)
    }

    unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
        off();
        let ret = self.heap
            .borrow_mut()
            .as_mut()
            .unwrap()
            .free(ptr, layout.size());
        on();
        ret
    }
}

/// Initialize the kernel heap
pub fn init(allocator: &mut KernelAllocator, start: usize, size: usize) {
    let heap = unsafe { Allocator::new(start as *mut u8, size) };
    let free_size = heap.size();

    allocator.set_heap(heap);

    bootlog!(
        "heap inited - start addr: 0x{:x}, end addr: 0x{:x}, {} bytes\n",
        start,
        start + size,
        free_size
    );
}
<|endoftext|>


<commit_before>#![feature(slice_concat_ext, custom_test_frameworks)]
// Custom test runner, to avoid libtest being wrapped around compiletest which wraps libtest.
#![test_runner(test_runner)]

use std::slice::SliceConcatExt;
use std::path::PathBuf;
use std::env;

use compiletest_rs as compiletest;
use colored::*;

fn miri_path() -> PathBuf {
    if rustc_test_suite().is_some() {
        PathBuf::from(option_env!("MIRI_PATH").unwrap())
    } else {
        PathBuf::from(concat!("target/", env!("PROFILE"), "/miri"))
    }
}

fn rustc_test_suite() -> Option<PathBuf> {
    option_env!("RUSTC_TEST_SUITE").map(PathBuf::from)
}

fn rustc_lib_path() -> PathBuf {
    option_env!("RUSTC_LIB_PATH").unwrap().into()
}

fn mk_config(mode: &str) -> compiletest::common::ConfigWithTemp {
    let mut config = compiletest::Config::default().tempdir();
    config.mode = mode.parse().expect("Invalid mode");
    config.rustc_path = miri_path();
    if rustc_test_suite().is_some() {
        config.run_lib_path = rustc_lib_path();
        config.compile_lib_path = rustc_lib_path();
    }
    config.filter = env::args().nth(1);
    config.host = get_host();
    config
}

fn compile_fail(path: &str, target: &str, opt: bool) {
    let opt_str = if opt { " with optimizations" } else { "" };
    eprintln!("{}", format!(
        "## Running compile-fail tests in {} against miri for target {}{}",
        path,
        target,
        opt_str
    ).green().bold());

    let mut flags = Vec::new();
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    flags.push("--edition 2018".to_owned());
    if opt {
        // Optimizing too aggressivley makes UB detection harder, but test at least
        // the default value.
        // FIXME: Opt level 3 ICEs during stack trace generation.
        flags.push("-Zmir-opt-level=1".to_owned());
    }

    let mut config = mk_config("compile-fail");
    config.src_base = PathBuf::from(path);
    config.target = target.to_owned();
    config.target_rustcflags = Some(flags.join(" "));
    compiletest::run_tests(&config);
}

fn miri_pass(path: &str, target: &str, opt: bool) {
    let opt_str = if opt { " with optimizations" } else { "" };
    eprintln!("{}", format!(
        "## Running run-pass tests in {} against miri for target {}{}",
        path,
        target,
        opt_str
    ).green().bold());

    let mut flags = Vec::new();
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    flags.push("--edition 2018".to_owned());
    if opt {
        // FIXME: We use opt level 1 because MIR inlining defeats the validation
        // whitelist.
        flags.push("-Zmir-opt-level=1".to_owned());
    }

    let mut config = mk_config("ui");
    config.src_base = PathBuf::from(path);
    config.target = target.to_owned();
    config.target_rustcflags = Some(flags.join(" "));
    compiletest::run_tests(&config);
}

/// Make sure the MIRI_SYSROOT env var is set
fn set_sysroot() {
    if std::env::var("MIRI_SYSROOT").is_ok() {
        // Nothing to do
        return;
    }
    let sysroot = std::process::Command::new("rustc")
        .arg("--print")
        .arg("sysroot")
        .output()
        .expect("rustc not found")
        .stdout;
    let sysroot = String::from_utf8(sysroot).expect("sysroot is not utf8");
    std::env::set_var("MIRI_SYSROOT", sysroot.trim());
}

fn get_host() -> String {
    let rustc = rustc_test_suite().unwrap_or(PathBuf::from("rustc"));
    let rustc_version = std::process::Command::new(rustc)
        .arg("-vV")
        .output()
        .expect("rustc not found for -vV")
        .stdout;
    let rustc_version = std::str::from_utf8(&rustc_version).expect("rustc -vV is not utf8");
    let version_meta = rustc_version::version_meta_for(&rustc_version)
        .expect("failed to parse rustc version info");
    version_meta.host
}

fn get_target() -> String {
    std::env::var("MIRI_TARGET").unwrap_or_else(|_| get_host())
}

fn run_pass_miri(opt: bool) {
    miri_pass("tests/run-pass", &get_target(), opt);
}

fn compile_fail_miri(opt: bool) {
    compile_fail("tests/compile-fail", &get_target(), opt);
}

fn test_runner(_tests: &[&()]) {
    set_sysroot();

    run_pass_miri(false);
    run_pass_miri(true);

    compile_fail_miri(false);
    compile_fail_miri(true);
}
<commit_msg>we have no whitelist any more, go back to opt-level 3<commit_after>#![feature(slice_concat_ext, custom_test_frameworks)]
// Custom test runner, to avoid libtest being wrapped around compiletest which wraps libtest.
#![test_runner(test_runner)]

use std::slice::SliceConcatExt;
use std::path::PathBuf;
use std::env;

use compiletest_rs as compiletest;
use colored::*;

fn miri_path() -> PathBuf {
    if rustc_test_suite().is_some() {
        PathBuf::from(option_env!("MIRI_PATH").unwrap())
    } else {
        PathBuf::from(concat!("target/", env!("PROFILE"), "/miri"))
    }
}

fn rustc_test_suite() -> Option<PathBuf> {
    option_env!("RUSTC_TEST_SUITE").map(PathBuf::from)
}

fn rustc_lib_path() -> PathBuf {
    option_env!("RUSTC_LIB_PATH").unwrap().into()
}

fn mk_config(mode: &str) -> compiletest::common::ConfigWithTemp {
    let mut config = compiletest::Config::default().tempdir();
    config.mode = mode.parse().expect("Invalid mode");
    config.rustc_path = miri_path();
    if rustc_test_suite().is_some() {
        config.run_lib_path = rustc_lib_path();
        config.compile_lib_path = rustc_lib_path();
    }
    config.filter = env::args().nth(1);
    config.host = get_host();
    config
}

fn compile_fail(path: &str, target: &str, opt: bool) {
    let opt_str = if opt { " with optimizations" } else { "" };
    eprintln!("{}", format!(
        "## Running compile-fail tests in {} against miri for target {}{}",
        path,
        target,
        opt_str
    ).green().bold());

    let mut flags = Vec::new();
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    flags.push("--edition 2018".to_owned());
    if opt {
        // Optimizing too aggressivley makes UB detection harder, but test at least
        // the default value.
        // FIXME: Opt level 3 ICEs during stack trace generation.
        flags.push("-Zmir-opt-level=1".to_owned());
    }

    let mut config = mk_config("compile-fail");
    config.src_base = PathBuf::from(path);
    config.target = target.to_owned();
    config.target_rustcflags = Some(flags.join(" "));
    compiletest::run_tests(&config);
}

fn miri_pass(path: &str, target: &str, opt: bool) {
    let opt_str = if opt { " with optimizations" } else { "" };
    eprintln!("{}", format!(
        "## Running run-pass tests in {} against miri for target {}{}",
        path,
        target,
        opt_str
    ).green().bold());

    let mut flags = Vec::new();
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    flags.push("--edition 2018".to_owned());
    if opt {
        flags.push("-Zmir-opt-level=3".to_owned());
    }

    let mut config = mk_config("ui");
    config.src_base = PathBuf::from(path);
    config.target = target.to_owned();
    config.target_rustcflags = Some(flags.join(" "));
    compiletest::run_tests(&config);
}

/// Make sure the MIRI_SYSROOT env var is set
fn set_sysroot() {
    if std::env::var("MIRI_SYSROOT").is_ok() {
        // Nothing to do
        return;
    }
    let sysroot = std::process::Command::new("rustc")
        .arg("--print")
        .arg("sysroot")
        .output()
        .expect("rustc not found")
        .stdout;
    let sysroot = String::from_utf8(sysroot).expect("sysroot is not utf8");
    std::env::set_var("MIRI_SYSROOT", sysroot.trim());
}

fn get_host() -> String {
    let rustc = rustc_test_suite().unwrap_or(PathBuf::from("rustc"));
    let rustc_version = std::process::Command::new(rustc)
        .arg("-vV")
        .output()
        .expect("rustc not found for -vV")
        .stdout;
    let rustc_version = std::str::from_utf8(&rustc_version).expect("rustc -vV is not utf8");
    let version_meta = rustc_version::version_meta_for(&rustc_version)
        .expect("failed to parse rustc version info");
    version_meta.host
}

fn get_target() -> String {
    std::env::var("MIRI_TARGET").unwrap_or_else(|_| get_host())
}

fn run_pass_miri(opt: bool) {
    miri_pass("tests/run-pass", &get_target(), opt);
}

fn compile_fail_miri(opt: bool) {
    compile_fail("tests/compile-fail", &get_target(), opt);
}

fn test_runner(_tests: &[&()]) {
    set_sysroot();

    run_pass_miri(false);
    run_pass_miri(true);

    compile_fail_miri(false);
    compile_fail_miri(true);
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use std::mem;
use rustyline;
use rustyline::error::ReadlineError as RLError;

#[derive(Debug)]
pub struct Readline {
    rl: rustyline::Editor<()>,
    prompt: String,
    buffer: String,
}

impl Readline {
    pub fn new(prompt: &str) -> Self {
        Readline {
            rl: rustyline::Editor::new(),
            prompt: prompt.to_owned(),
            buffer: String::with_capacity(128),
        }
    }
}

impl Read for Readline {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.push_str(&line);
                    self.buffer.push('\n');
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Eof) => return Ok(0),
                Err(RLError::Interrupted) => return Ok(0),
                #[cfg(unix)]
                Err(RLError::Errno(num)) => return Err(io::Error::from_raw_os_error(num)),
                #[cfg(windows)]
                Err(RLError::WindowResize) => return Err(io::ErrorKind::Other.into()),
                #[cfg(windows)]
                Err(RLError::Decode(_)) => return Err(io::ErrorKind::InvalidData.into()),
            };
        }

        // if buf > buffer, drain entire buffer
        if buf.len() >= self.buffer.len() {
            let result = buf.write(&self.buffer.as_bytes());
            self.buffer.clear();
            result
        } else {
            let mut copy = String::with_capacity(128);
            mem::swap(&mut self.buffer, &mut copy);
            buf.write(&copy.into_bytes())
        }
    }
}

impl io::BufRead for Readline {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.push_str(&line);
                    self.buffer.push('\n');
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Eof) => (),
                Err(RLError::Interrupted) => (),
                #[cfg(unix)]
                Err(RLError::Errno(num)) => return Err(io::Error::from_raw_os_error(num)),
                #[cfg(windows)]
                Err(RLError::WindowResize) => return Err(io::ErrorKind::Other.into()),
                #[cfg(windows)]
                Err(RLError::Decode(_)) => return Err(io::ErrorKind::InvalidData.into()),
            }
        }

        Ok(&self.buffer.as_bytes())
    }

    fn consume(&mut self, amt: usize) {
        self.buffer.drain(0..amt);
    }
}
<commit_msg>Fix readline on *nix platforms<commit_after>use std::io;
use std::io::prelude::*;
use std::mem;
use rustyline;
use rustyline::error::ReadlineError as RLError;

#[derive(Debug)]
pub struct Readline {
    rl: rustyline::Editor<()>,
    prompt: String,
    buffer: String,
}

impl Readline {
    pub fn new(prompt: &str) -> Self {
        Readline {
            rl: rustyline::Editor::new(),
            prompt: prompt.to_owned(),
            buffer: String::with_capacity(128),
        }
    }
}

impl Read for Readline {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.push_str(&line);
                    self.buffer.push('\n');
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Eof) => return Ok(0),
                Err(RLError::Interrupted) => return Ok(0),
                #[cfg(unix)]
                Err(RLError::Char(_)) => return Ok(0),
                #[cfg(unix)]
                Err(RLError::Errno(err)) => return Err(io::Error::from(err)),
                #[cfg(windows)]
                Err(RLError::WindowResize) => return Err(io::ErrorKind::Other.into()),
                #[cfg(windows)]
                Err(RLError::Decode(_)) => return Err(io::ErrorKind::InvalidData.into()),
            };
        }

        // if buf > buffer, drain entire buffer
        if buf.len() >= self.buffer.len() {
            let result = buf.write(&self.buffer.as_bytes());
            self.buffer.clear();
            result
        } else {
            let mut copy = String::with_capacity(128);
            mem::swap(&mut self.buffer, &mut copy);
            buf.write(&copy.into_bytes())
        }
    }
}

impl io::BufRead for Readline {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if self.buffer.is_empty() {
            match self.rl.readline(&self.prompt) {
                Ok(line) => {
                    self.rl.add_history_entry(&line);
                    self.buffer.push_str(&line);
                    self.buffer.push('\n');
                },
                Err(RLError::Io(err)) => return Err(err),
                Err(RLError::Eof) => (),
                Err(RLError::Interrupted) => (),
                #[cfg(unix)]
                Err(RLError::Char(_)) => (),
                #[cfg(unix)]
                Err(RLError::Errno(err)) => return Err(io::Error::from(err)),
                #[cfg(windows)]
                Err(RLError::WindowResize) => return Err(io::ErrorKind::Other.into()),
                #[cfg(windows)]
                Err(RLError::Decode(_)) => return Err(io::ErrorKind::InvalidData.into()),
            }
        }

        Ok(&self.buffer.as_bytes())
    }

    fn consume(&mut self, amt: usize) {
        self.buffer.drain(0..amt);
    }
}
<|endoftext|>


<commit_before>use bits;
use data;
use voice;

pub trait Storage {
    type Input: Copy;
    type Buf;

    fn size(&self) -> usize;
    fn buf(&mut self) -> &mut Self::Buf;
    fn add(&mut self, item: Self::Input, pos: usize);
    fn reset(&mut self) {}
}

macro_rules! storage_type {
    ($name: ident, [$input: ty; $size: expr]) => {
        pub struct $name([$input; $size]);

        impl $name {
            pub fn new() -> Self { $name([Default::default(); $size]) }
        }

        impl Storage for $name {
            type Input = $input;
            type Buf = [$input; $size];

            fn size(&self) -> usize { $size }
            fn buf(&mut self) -> &mut Self::Buf { &mut self.0 }
            fn add(&mut self, item: Self::Input, pos: usize) { self.0[pos] = item; }
        }
    };
}

storage_type!(VoiceHeaderStorage, [bits::Hexbit; 36]);
storage_type!(VoiceFrameStorage, [bits::Dibit; voice::consts::FRAME_DIBITS]);
storage_type!(VoiceExtraStorage, [bits::Hexbit; 24]);
storage_type!(DataHeaderStorage, [bits::Dibit; 48]);
storage_type!(DataPayloadStorage, [bits::Dibit; data::consts::CODING_DIBITS]);

pub struct DibitStorage {
    buf: u64,
    size: usize,
}

impl DibitStorage {
    pub fn new(size: usize) -> DibitStorage {
        assert!(size <= 32);

        DibitStorage {
            buf: 0,
            size: size,
        }
    }
}

impl Storage for DibitStorage {
    type Input = bits::Dibit;
    type Buf = u64;

    fn size(&self) -> usize { self.size }
    fn buf(&mut self) -> &mut u64 { &mut self.buf }

    fn add(&mut self, item: Self::Input, _: usize) {
        self.buf <<= 2;
        self.buf |= item.bits() as u64;
    }

    fn reset(&mut self) { self.buf = 0; }
}

pub struct Buffer<S: Storage> {
    storage: S,
    pos: usize,
}

impl<S: Storage> Buffer<S> {
    pub fn new(storage: S) -> Buffer<S> {
        Buffer {
            storage: storage,
            pos: 0,
        }
    }

    pub fn reset(&mut self) { self.pos = 0; }

    /// the buffer is reset once `Some` is returned.
    pub fn feed(&mut self, item: S::Input) -> Option<&mut S::Buf> {
        if self.pos == 0 {
            self.storage.reset();
        }

        self.storage.add(item, self.pos);
        self.pos += 1;

        if self.pos == self.storage.size() {
            self.reset();
            Some(self.storage.buf())
        } else {
            None
        }
    }
}
<commit_msg>Remove DataHeaderStorage for now<commit_after>use bits;
use data;
use voice;

pub trait Storage {
    type Input: Copy;
    type Buf;

    fn size(&self) -> usize;
    fn buf(&mut self) -> &mut Self::Buf;
    fn add(&mut self, item: Self::Input, pos: usize);
    fn reset(&mut self) {}
}

macro_rules! storage_type {
    ($name: ident, [$input: ty; $size: expr]) => {
        pub struct $name([$input; $size]);

        impl $name {
            pub fn new() -> Self { $name([Default::default(); $size]) }
        }

        impl Storage for $name {
            type Input = $input;
            type Buf = [$input; $size];

            fn size(&self) -> usize { $size }
            fn buf(&mut self) -> &mut Self::Buf { &mut self.0 }
            fn add(&mut self, item: Self::Input, pos: usize) { self.0[pos] = item; }
        }
    };
}

storage_type!(VoiceHeaderStorage, [bits::Hexbit; 36]);
storage_type!(VoiceFrameStorage, [bits::Dibit; voice::consts::FRAME_DIBITS]);
storage_type!(VoiceExtraStorage, [bits::Hexbit; 24]);
storage_type!(DataPayloadStorage, [bits::Dibit; data::consts::CODING_DIBITS]);

pub struct DibitStorage {
    buf: u64,
    size: usize,
}

impl DibitStorage {
    pub fn new(size: usize) -> DibitStorage {
        assert!(size <= 32);

        DibitStorage {
            buf: 0,
            size: size,
        }
    }
}

impl Storage for DibitStorage {
    type Input = bits::Dibit;
    type Buf = u64;

    fn size(&self) -> usize { self.size }
    fn buf(&mut self) -> &mut u64 { &mut self.buf }

    fn add(&mut self, item: Self::Input, _: usize) {
        self.buf <<= 2;
        self.buf |= item.bits() as u64;
    }

    fn reset(&mut self) { self.buf = 0; }
}

pub struct Buffer<S: Storage> {
    storage: S,
    pos: usize,
}

impl<S: Storage> Buffer<S> {
    pub fn new(storage: S) -> Buffer<S> {
        Buffer {
            storage: storage,
            pos: 0,
        }
    }

    pub fn reset(&mut self) { self.pos = 0; }

    /// the buffer is reset once `Some` is returned.
    pub fn feed(&mut self, item: S::Input) -> Option<&mut S::Buf> {
        if self.pos == 0 {
            self.storage.reset();
        }

        self.storage.add(item, self.pos);
        self.pos += 1;

        if self.pos == self.storage.size() {
            self.reset();
            Some(self.storage.buf())
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_s3::{config, Client, Error, Region, RetryConfig, PKG_VERSION};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// The number of (re)tries.
    #[structopt(short, long, default_value = "2")]
    tries: u32,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// Shows your buckets.
async fn show_num_buckets(client: &Client) -> Result<(), Error> {
    let resp = client.list_buckets().send().await?;
    let buckets = resp.buckets().unwrap_or_default();

    println!("Found {} buckets in all regions.", buckets.len());

    Ok(())
}

/// Displays how many Amazon S3 buckets you have.
/// # Arguments
///
/// * `[-r REGION]` - The Region in which the client is created.
///   If not supplied, uses the value of the **AWS_REGION** environment variable.
///   If the environment variable is not set, defaults to **us-west-2**.
/// * `[-t TRIES]` - The number of times to (re)try the request.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();

    let Opt {
        region,
        tries,
        verbose,
    } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));
    println!();

    if verbose {
        println!("S3 client version: {}", PKG_VERSION);
        println!(
            "Region:            {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!("Retries:           {}", &tries);
        println!();
    }

    assert_ne!(tries, 0, "You cannot set zero retries.");

    // snippet-start:[custom_retries.rust.set_retries]
    let shared_config = aws_config::from_env().region(region_provider).load().await;

    // Construct an S3 client with customized retry configuration.
    let client = Client::from_conf(
        // Start with the shared environment configuration.
        config::Builder::from(&shared_config)
            // Set max attempts.
            // If tries is 1, there are no retries.
            .retry_config(RetryConfig::new().with_max_attempts(tries))
            .build(),
    );
    // snippet-end:[custom_retries.rust.set_retries]

    show_num_buckets(&client).await
}
<commit_msg>Fix retry config example after breaking change in the Rust SDK<commit_after>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_s3::{config, Client, Error, Region, RetryConfig, PKG_VERSION};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// The number of (re)tries.
    #[structopt(short, long, default_value = "2")]
    tries: u32,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// Shows your buckets.
async fn show_num_buckets(client: &Client) -> Result<(), Error> {
    let resp = client.list_buckets().send().await?;
    let buckets = resp.buckets().unwrap_or_default();

    println!("Found {} buckets in all regions.", buckets.len());

    Ok(())
}

/// Displays how many Amazon S3 buckets you have.
/// # Arguments
///
/// * `[-r REGION]` - The Region in which the client is created.
///   If not supplied, uses the value of the **AWS_REGION** environment variable.
///   If the environment variable is not set, defaults to **us-west-2**.
/// * `[-t TRIES]` - The number of times to (re)try the request.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();

    let Opt {
        region,
        tries,
        verbose,
    } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));
    println!();

    if verbose {
        println!("S3 client version: {}", PKG_VERSION);
        println!(
            "Region:            {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!("Retries:           {}", &tries);
        println!();
    }

    assert_ne!(tries, 0, "You cannot set zero retries.");

    // snippet-start:[custom_retries.rust.set_retries]
    let shared_config = aws_config::from_env().region(region_provider).load().await;

    // Construct an S3 client with customized retry configuration.
    let client = Client::from_conf(
        // Start with the shared environment configuration.
        config::Builder::from(&shared_config)
            // Set max attempts.
            // If tries is 1, there are no retries.
            .retry_config(RetryConfig::standard().with_max_attempts(tries))
            .build(),
    );
    // snippet-end:[custom_retries.rust.set_retries]

    show_num_buckets(&client).await
}
<|endoftext|>


<commit_before>use std::collections::hashmap::HashMap;
use time;
use std::mem;
use sync::comm::Receiver;
use image::{GenericImage, SubImage};
use image::{Pixel, ImageBuf, Rgba};
use image;
use glfw;
use glfw::Context as _Context;
use gfx;
use util;

static FONT_DATA: &'static [u8] = include_bin!("../assets/font.png");

pub struct Canvas {
    title: String,
    dim: [u32, ..2],
    frame_interval: Option<f64>,
    image_collector: ImageCollector,
    font_glyphs: HashMap<char, Image>,
}

/// Toplevel graphics drawing and input reading context.
impl Canvas {
    pub fn new() -> Canvas {
        let mut ret = Canvas {
            title: "window".to_string(),
            dim: [640, 360],
            frame_interval: None,
            image_collector: ImageCollector::new(),
            font_glyphs: HashMap::new(),
        };
        ret.init_font();
        ret
    }

    pub fn set_title(&mut self, title: &str) -> &mut Canvas {
        self.title = title.to_string();
        self
    }

    /// Set the frame rate.
    pub fn set_frame_interval(&mut self, interval_s: f64) -> &mut Canvas {
        assert!(interval_s > 0.00001);
        self.frame_interval = Some(interval_s);
        self
    }

    /// Set the resolution.
    pub fn set_dim(&mut self, dim: [u32, ..2]) -> &mut Canvas {
        self.dim = dim;
        self
    }

    pub fn add_image<P: Pixel<u8>, I: GenericImage<P>>(
        &mut self, image: I) -> Image {
        Image(self.image_collector.push(image))
    }

    /// Start running the engine, return an event iteration.
    pub fn run(&mut self) -> Context {
        Context::new(
            self.dim,
            self.title.as_slice(),
            Some(1.0 / 30.0),
            self.image_collector.build_atlas())
    }

    fn init_font(&mut self) {
        let mut font_sheet = image::load_from_memory(FONT_DATA, image::PNG).unwrap();
        for i in range(0u32, 96u32) {
            let x = 8u32 * (i % 16u32);
            let y = 8u32 * (i / 16u32);
            let glyph = self.add_image(SubImage::new(&mut font_sheet, x, y, 8, 8));
            self.font_glyphs.insert((i + 32) as u8 as char, glyph);
        }
    }
}

pub struct Context {
    glfw: glfw::Glfw,
    window: glfw::Window,
    events: Receiver<(f64, glfw::WindowEvent)>,
    graphics: gfx::Graphics<gfx::GlDevice, gfx::GlCommandBuffer>,
    frame: gfx::Frame,

    state: State,
    frame_interval: Option<f64>,
    last_render_time: f64,
    atlas: Atlas,
}

#[deriving(PartialEq)]
enum State {
    Normal,
    EndFrame,
}

impl Context {
    fn new(
        dim: [u32, ..2],
        title: &str,
        frame_interval: Option<f64>,
        atlas: Atlas) -> Context {

        let glfw = glfw::init(glfw::FAIL_ON_ERRORS).unwrap();
        let (window, events) = glfw
            .create_window(dim[0], dim[1],
            title.as_slice(), glfw::Windowed)
            .expect("Failed to open window");
        window.make_current();
        glfw.set_error_callback(glfw::FAIL_ON_ERRORS);
        window.set_key_polling(true);
        window.set_char_polling(true);

        let device = gfx::GlDevice::new(|s| window.get_proc_address(s));
        let graphics = gfx::Graphics::new(device);
        let frame = gfx::Frame::new(dim[0] as u16, dim[1] as u16);

        Context {
            glfw: glfw,
            window: window,
            events: events,
            graphics: graphics,
            frame: frame,
            state: Normal,
            frame_interval: frame_interval,
            last_render_time: time::precise_time_s(),
            atlas: atlas,
        }
    }

    /// Clear the screen
    pub fn clear(&mut self, color: [f32, ..4]) {
        self.graphics.clear(
            gfx::ClearData {
                color: color,
                depth: 1.0,
                stencil: 0,
            }, gfx::Color, &self.frame);
    }

    pub fn draw_image(&mut self, offset: [int, ..2], image: Image) {
        unimplemented!();
    }
}

pub enum Event<'a> {
    Render(&'a mut Context),
    Input(int), // TODO: Proper input type
}

impl<'a> Iterator<Event<'a>> for Context {
    fn next(&mut self) -> Option<Event<'a>> {
        // After a render event, control will return here on a new
        // iter call. Do post-render work here.
        if self.state == EndFrame {
            self.state = Normal;
            self.graphics.end_frame();
            self.window.swap_buffers();
        }

        loop {
            if self.window.should_close() {
                return None;
            }

            self.glfw.poll_events();

            match self.events.try_recv() {
                Ok(_event) => {
                    // TODO: Process event.
                    return Some(Input(123))
                }
                _ => ()
            }

            let t = time::precise_time_s();
            if self.frame_interval.map_or(true,
                |x| t - self.last_render_time >= x) {
                self.last_render_time = t;

                // Time to render, must return a handle to self.
                // XXX: Need unsafe hackery to get around lifetimes check.
                self.state = EndFrame;
                unsafe {
                    return Some(Render(mem::transmute(self)))
                }
            }
        }
    }
}

/// Drawable images stored in the Canvas.
#[deriving(Clone, PartialEq)]
pub struct Image(uint);

struct ImageCollector {
    images: Vec<ImageBuf<Rgba<u8>>>,
    offsets: Vec<[u32, ..2]>,
}

impl ImageCollector {
    fn new() -> ImageCollector {
        ImageCollector {
            images: vec![],
            offsets: vec![],
        }
    }

    fn push<P: Pixel<u8>, I: GenericImage<P>>(
        &mut self, mut image: I) -> uint {
        let (pos, dim) = util::crop_alpha(&image);
        let cropped = SubImage::new(&mut image, pos[0], pos[1], dim[0], dim[1]);

        let (w, h) = cropped.dimensions();
        let img = ImageBuf::from_pixels(
            cropped.pixels().map::<Rgba<u8>>(
                |(_x, _y, p)| p.to_rgba())
            .collect(),
            w, h);
        self.images.push(img);
        self.offsets.push(pos);
        self.images.len()
    }

    fn build_atlas(&mut self) -> Atlas {
        Atlas::new(&self.images, &self.offsets)
    }

    fn rendering_offset(&self, idx: uint) -> [u32, ..2] { self.offsets[idx] }
}

struct Atlas {
    image: ImageBuf<Rgba<u8>>,
    bounds: Vec<([u32, ..2], [u32, ..2])>,
    offsets: Vec<[u32, ..2]>,
}

impl Atlas {
    fn new(images: &Vec<ImageBuf<Rgba<u8>>>, offsets: &Vec<[u32, ..2]>) -> Atlas {
        let (image, bounds) = util::build_atlas(images);

        // TODO: Replace with offsets.clone() when Rust supports fixed size array cloning.
        let mut offsets_clone = vec![]; for i in offsets.iter() { offsets_clone.push(*i); }

        Atlas {
            image: image,
            bounds: bounds,
            offsets: offsets_clone,
        }
    }
}
<commit_msg>Non-reference fluent interface<commit_after>use std::collections::hashmap::HashMap;
use time;
use std::mem;
use sync::comm::Receiver;
use image::{GenericImage, SubImage};
use image::{Pixel, ImageBuf, Rgba};
use image;
use glfw;
use glfw::Context as _Context;
use gfx;
use util;

static FONT_DATA: &'static [u8] = include_bin!("../assets/font.png");

pub struct Canvas {
    title: String,
    dim: [u32, ..2],
    frame_interval: Option<f64>,
    image_collector: ImageCollector,
    font_glyphs: HashMap<char, Image>,
}

/// Toplevel graphics drawing and input reading context.
impl Canvas {
    pub fn new() -> Canvas {
        let mut ret = Canvas {
            title: "window".to_string(),
            dim: [640, 360],
            frame_interval: None,
            image_collector: ImageCollector::new(),
            font_glyphs: HashMap::new(),
        };
        ret.init_font();
        ret
    }

    pub fn set_title(mut self, title: &str) -> Canvas {
        self.title = title.to_string();
        self
    }

    /// Set the frame rate.
    pub fn set_frame_interval(mut self, interval_s: f64) -> Canvas {
        assert!(interval_s > 0.00001);
        self.frame_interval = Some(interval_s);
        self
    }

    /// Set the resolution.
    pub fn set_dim(mut self, dim: [u32, ..2]) -> Canvas {
        self.dim = dim;
        self
    }

    pub fn add_image<P: Pixel<u8>, I: GenericImage<P>>(
        &mut self, image: I) -> Image {
        Image(self.image_collector.push(image))
    }

    /// Start running the engine, return an event iteration.
    pub fn run(&mut self) -> Context {
        Context::new(
            self.dim,
            self.title.as_slice(),
            Some(1.0 / 30.0),
            self.image_collector.build_atlas())
    }

    fn init_font(&mut self) {
        let mut font_sheet = image::load_from_memory(FONT_DATA, image::PNG).unwrap();
        for i in range(0u32, 96u32) {
            let x = 8u32 * (i % 16u32);
            let y = 8u32 * (i / 16u32);
            let glyph = self.add_image(SubImage::new(&mut font_sheet, x, y, 8, 8));
            self.font_glyphs.insert((i + 32) as u8 as char, glyph);
        }
    }
}

pub struct Context {
    glfw: glfw::Glfw,
    window: glfw::Window,
    events: Receiver<(f64, glfw::WindowEvent)>,
    graphics: gfx::Graphics<gfx::GlDevice, gfx::GlCommandBuffer>,
    frame: gfx::Frame,

    state: State,
    frame_interval: Option<f64>,
    last_render_time: f64,
    atlas: Atlas,
}

#[deriving(PartialEq)]
enum State {
    Normal,
    EndFrame,
}

impl Context {
    fn new(
        dim: [u32, ..2],
        title: &str,
        frame_interval: Option<f64>,
        atlas: Atlas) -> Context {

        let glfw = glfw::init(glfw::FAIL_ON_ERRORS).unwrap();
        let (window, events) = glfw
            .create_window(dim[0], dim[1],
            title.as_slice(), glfw::Windowed)
            .expect("Failed to open window");
        window.make_current();
        glfw.set_error_callback(glfw::FAIL_ON_ERRORS);
        window.set_key_polling(true);
        window.set_char_polling(true);

        let device = gfx::GlDevice::new(|s| window.get_proc_address(s));
        let graphics = gfx::Graphics::new(device);
        let frame = gfx::Frame::new(dim[0] as u16, dim[1] as u16);

        Context {
            glfw: glfw,
            window: window,
            events: events,
            graphics: graphics,
            frame: frame,
            state: Normal,
            frame_interval: frame_interval,
            last_render_time: time::precise_time_s(),
            atlas: atlas,
        }
    }

    /// Clear the screen
    pub fn clear(&mut self, color: [f32, ..4]) {
        self.graphics.clear(
            gfx::ClearData {
                color: color,
                depth: 1.0,
                stencil: 0,
            }, gfx::Color, &self.frame);
    }

    pub fn draw_image(&mut self, offset: [int, ..2], image: Image) {
        unimplemented!();
    }
}

pub enum Event<'a> {
    Render(&'a mut Context),
    Input(int), // TODO: Proper input type
}

impl<'a> Iterator<Event<'a>> for Context {
    fn next(&mut self) -> Option<Event<'a>> {
        // After a render event, control will return here on a new
        // iter call. Do post-render work here.
        if self.state == EndFrame {
            self.state = Normal;
            self.graphics.end_frame();
            self.window.swap_buffers();
        }

        loop {
            if self.window.should_close() {
                return None;
            }

            self.glfw.poll_events();

            match self.events.try_recv() {
                Ok(_event) => {
                    // TODO: Process event.
                    return Some(Input(123))
                }
                _ => ()
            }

            let t = time::precise_time_s();
            if self.frame_interval.map_or(true,
                |x| t - self.last_render_time >= x) {
                self.last_render_time = t;

                // Time to render, must return a handle to self.
                // XXX: Need unsafe hackery to get around lifetimes check.
                self.state = EndFrame;
                unsafe {
                    return Some(Render(mem::transmute(self)))
                }
            }
        }
    }
}

/// Drawable images stored in the Canvas.
#[deriving(Clone, PartialEq)]
pub struct Image(uint);

struct ImageCollector {
    images: Vec<ImageBuf<Rgba<u8>>>,
    offsets: Vec<[u32, ..2]>,
}

impl ImageCollector {
    fn new() -> ImageCollector {
        ImageCollector {
            images: vec![],
            offsets: vec![],
        }
    }

    fn push<P: Pixel<u8>, I: GenericImage<P>>(
        &mut self, mut image: I) -> uint {
        let (pos, dim) = util::crop_alpha(&image);
        let cropped = SubImage::new(&mut image, pos[0], pos[1], dim[0], dim[1]);

        let (w, h) = cropped.dimensions();
        let img = ImageBuf::from_pixels(
            cropped.pixels().map::<Rgba<u8>>(
                |(_x, _y, p)| p.to_rgba())
            .collect(),
            w, h);
        self.images.push(img);
        self.offsets.push(pos);
        self.images.len()
    }

    fn build_atlas(&mut self) -> Atlas {
        Atlas::new(&self.images, &self.offsets)
    }

    fn rendering_offset(&self, idx: uint) -> [u32, ..2] { self.offsets[idx] }
}

struct Atlas {
    image: ImageBuf<Rgba<u8>>,
    bounds: Vec<([u32, ..2], [u32, ..2])>,
    offsets: Vec<[u32, ..2]>,
}

impl Atlas {
    fn new(images: &Vec<ImageBuf<Rgba<u8>>>, offsets: &Vec<[u32, ..2]>) -> Atlas {
        let (image, bounds) = util::build_atlas(images);

        // TODO: Replace with offsets.clone() when Rust supports fixed size array cloning.
        let mut offsets_clone = vec![]; for i in offsets.iter() { offsets_clone.push(*i); }

        Atlas {
            image: image,
            bounds: bounds,
            offsets: offsets_clone,
        }
    }
}
<|endoftext|>


<commit_before>
use tokio::runtime::Runtime;
use log::warn;
use hyper::{Uri, Body, Request};
use hyper::header::CONTENT_TYPE;
use hyper::client::{Client, HttpConnector};
use hyper_tls::HttpsConnector;

use crate::Notice;
use crate::AirbrakeConfig;

pub struct AirbrakeClient {
    client: Client<HttpsConnector<HttpConnector>>,
    config: AirbrakeConfig
}

impl AirbrakeClient {
    pub fn new(config: AirbrakeConfig) -> AirbrakeClient {
        let connector = HttpsConnector::new();
        let client = Client::builder().build(connector);

        AirbrakeClient {
            client: client,
            config: config
        }
    }

    fn request<T>(&self, uri: Uri, payload: T) -> Request<Body>
    where T: Into<Body>
    {
        Request::post(uri)
            .header(CONTENT_TYPE, "application/json")
            .body(payload.into())
            .expect("Request creation failed unexpectedly")
    }

    async fn send(&self, request: Request<Body>) -> () {
        let response = self.client.request(request).await;
        match response {
            Ok ( response ) => (),
            Err ( x ) => warn!("notification failed")
        }
    }

    pub fn notify(&self, notice: Notice) {
        let endpoint = self.config.endpoint_uri();
        let request = self.request(endpoint, notice);

        let mut runtime = Runtime::new().unwrap();
        runtime.block_on(self.send(request));
    }
}
<commit_msg>rename internal function names for clarity<commit_after>
use tokio::runtime::Runtime;
use log::warn;
use hyper::{Uri, Body, Request};
use hyper::header::CONTENT_TYPE;
use hyper::client::{Client, HttpConnector};
use hyper_tls::HttpsConnector;

use crate::Notice;
use crate::AirbrakeConfig;

pub struct AirbrakeClient {
    client: Client<HttpsConnector<HttpConnector>>,
    config: AirbrakeConfig
}

impl AirbrakeClient {
    pub fn new(config: AirbrakeConfig) -> AirbrakeClient {
        let connector = HttpsConnector::new();
        let client = Client::builder().build(connector);

        AirbrakeClient {
            client: client,
            config: config
        }
    }

    fn build_request<T>(&self, uri: Uri, payload: T) -> Request<Body>
    where T: Into<Body>
    {
        Request::post(uri)
            .header(CONTENT_TYPE, "application/json")
            .body(payload.into())
            .expect("Request creation failed unexpectedly")
    }

    async fn send_request(&self, request: Request<Body>) -> () {
        let response = self.client.request(request).await;
        match response {
            Ok ( response ) => (),
            Err ( x ) => warn!("notification failed")
        }
    }

    pub fn notify(&self, notice: Notice) {
        let endpoint = self.config.endpoint_uri();
        let request = self.build_request(endpoint, notice);

        let mut runtime = Runtime::new().unwrap();
        runtime.block_on(self.send_request(request));
    }
}
<|endoftext|>


<commit_before>use std::net::{UdpSocket, SocketAddr};
use std::io::Error;
use std::str::FromStr;

extern crate clock_ticks;
extern crate rand;

/// Client socket for statsd servers.
///
/// After creating a metric you can use `Client`
/// to send metrics to the configured statsd server
///
/// # Example
///
/// Creating a client and sending metrics is easy.
///
/// ```ignore
/// use statsd::client::Client;
///
/// let client = Client::new("127.0.0.1:8125", "myapp.");
/// client.incr("some.metric.completed");
/// ```
pub struct Client {
    socket: UdpSocket,
    server_address: SocketAddr,
    prefix: String
}

impl Client {
    /// Construct a new statsd client given an host/port & prefix
    pub fn new(host: &str, prefix: &str) -> Result<Client, Error> {
        // Bind to a generic port as we'll only be writing on this
        // socket.
        let client_address = SocketAddr::from_str("0.0.0.0:0").unwrap();
        let mut socket = try!(UdpSocket::bind(client_address));

        let server_address = SocketAddr::from_str(host).ok().unwrap();
        Ok(Client {
            socket: socket,
            prefix: prefix.to_string(),
            server_address: server_address
        })
    }

    /// Increment a metric by 1
    ///
    /// ```ignore
    /// # Increment a given metric by 1.
    /// client.incr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn incr(&mut self, metric: &str) {
        self.count(metric, 1.0);
    }

    /// Decrement a metric by -1
    ///
    /// ```ignore
    /// # Decrement a given metric by 1
    /// client.decr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn decr(&mut self, metric: &str) {
        self.count(metric, -1.0);
    }

    /// Modify a counter by `value`.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a sampling rate of 1.0.
    ///
    /// ```ignore
    /// # Increment by 12
    /// client.count("metric.completed", 12);
    /// ```
    pub fn count(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Modify a counter by `value` only x% of the time.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a custom sampling rate.
    ///
    ///
    /// ```ignore
    /// # Increment by 4 50% of the time.
    /// client.sampled_count("metric.completed", 4, 0.5);
    /// ```
    pub fn sampled_count(&mut self, metric: &str, value: f64, rate: f64) {
        if rand::random::<f64>() < rate {
            return
        }
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Set a gauge value.
    ///
    /// ```ignore
    /// # set a gauge to 9001
    /// client.gauge("power_level.observed", 9001);
    /// ```
    pub fn gauge(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|g", self.prefix, metric, value);
        self.send(data);
    }

    /// Send a timer value.
    ///
    /// The value is expected to be in ms.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.timer("response.duration", 10.123);
    /// ```
    pub fn timer(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|ms", self.prefix, metric, value);
        self.send(data);
    }

    /// Time a block of code.
    ///
    /// The passed closure will be timed and executed. The block's
    /// duration will be sent as a metric.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.time("response.duration", || {
    ///   # Your code here.
    /// });
    /// ```
    pub fn time<F>(&mut self, metric: &str, callable: F) where F : Fn() {
        let start = clock_ticks::precise_time_ms();
        callable();
        let end = clock_ticks::precise_time_ms();
        let data = format!("{}.{}:{}|ms", self.prefix, metric, end - start);
        self.send(data);
    }

    /// Send data along the UDP socket.
    fn send(&mut self, data: String) {
        let _ = self.socket.send_to(data.as_bytes(), self.server_address);
    }
}
<commit_msg>Remove unneeded mut ref.<commit_after>use std::net::{UdpSocket, SocketAddr};
use std::io::Error;
use std::str::FromStr;

extern crate clock_ticks;
extern crate rand;

/// Client socket for statsd servers.
///
/// After creating a metric you can use `Client`
/// to send metrics to the configured statsd server
///
/// # Example
///
/// Creating a client and sending metrics is easy.
///
/// ```ignore
/// use statsd::client::Client;
///
/// let client = Client::new("127.0.0.1:8125", "myapp.");
/// client.incr("some.metric.completed");
/// ```
pub struct Client {
    socket: UdpSocket,
    server_address: SocketAddr,
    prefix: String
}

impl Client {
    /// Construct a new statsd client given an host/port & prefix
    pub fn new(host: &str, prefix: &str) -> Result<Client, Error> {
        // Bind to a generic port as we'll only be writing on this
        // socket.
        let client_address = SocketAddr::from_str("0.0.0.0:0").unwrap();
        let socket = try!(UdpSocket::bind(client_address));

        let server_address = SocketAddr::from_str(host).ok().unwrap();
        Ok(Client {
            socket: socket,
            prefix: prefix.to_string(),
            server_address: server_address
        })
    }

    /// Increment a metric by 1
    ///
    /// ```ignore
    /// # Increment a given metric by 1.
    /// client.incr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn incr(&mut self, metric: &str) {
        self.count(metric, 1.0);
    }

    /// Decrement a metric by -1
    ///
    /// ```ignore
    /// # Decrement a given metric by 1
    /// client.decr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn decr(&mut self, metric: &str) {
        self.count(metric, -1.0);
    }

    /// Modify a counter by `value`.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a sampling rate of 1.0.
    ///
    /// ```ignore
    /// # Increment by 12
    /// client.count("metric.completed", 12);
    /// ```
    pub fn count(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Modify a counter by `value` only x% of the time.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a custom sampling rate.
    ///
    ///
    /// ```ignore
    /// # Increment by 4 50% of the time.
    /// client.sampled_count("metric.completed", 4, 0.5);
    /// ```
    pub fn sampled_count(&mut self, metric: &str, value: f64, rate: f64) {
        if rand::random::<f64>() < rate {
            return
        }
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Set a gauge value.
    ///
    /// ```ignore
    /// # set a gauge to 9001
    /// client.gauge("power_level.observed", 9001);
    /// ```
    pub fn gauge(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|g", self.prefix, metric, value);
        self.send(data);
    }

    /// Send a timer value.
    ///
    /// The value is expected to be in ms.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.timer("response.duration", 10.123);
    /// ```
    pub fn timer(&mut self, metric: &str, value: f64) {
        let data = format!("{}.{}:{}|ms", self.prefix, metric, value);
        self.send(data);
    }

    /// Time a block of code.
    ///
    /// The passed closure will be timed and executed. The block's
    /// duration will be sent as a metric.
    ///
    /// ```ignore
    /// # pass a duration value
    /// client.time("response.duration", || {
    ///   # Your code here.
    /// });
    /// ```
    pub fn time<F>(&mut self, metric: &str, callable: F) where F : Fn() {
        let start = clock_ticks::precise_time_ms();
        callable();
        let end = clock_ticks::precise_time_ms();
        let data = format!("{}.{}:{}|ms", self.prefix, metric, end - start);
        self.send(data);
    }

    /// Send data along the UDP socket.
    fn send(&mut self, data: String) {
        let _ = self.socket.send_to(data.as_bytes(), self.server_address);
    }
}
<|endoftext|>