<commit_before>use std::str;
use std::collections::HashMap;

use rocksdb::{Writable, IteratorMode, Direction, WriteBatch};
use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use RocksDBIndexStore;
use errors::{RocksDBReadError, RocksDBWriteError};
use document_index::DocRef;
use key_builder::KeyBuilder;
use search::doc_id_set::DocIdSet;

#[derive(Debug)]
pub enum SegmentMergeError {
    TooManyDocs,
    RocksDBReadError(RocksDBReadError),
    RocksDBWriteError(RocksDBWriteError),
}

impl From<RocksDBReadError> for SegmentMergeError {
    fn from(e: RocksDBReadError) -> SegmentMergeError {
        SegmentMergeError::RocksDBReadError(e)
    }
}

impl From<RocksDBWriteError> for SegmentMergeError {
    fn from(e: RocksDBWriteError) -> SegmentMergeError {
        SegmentMergeError::RocksDBWriteError(e)
    }
}

impl From<SegmentMergeError> for String {
    fn from(e: SegmentMergeError) -> String {
        format!("{:?}", e).to_string()
    }
}

impl RocksDBIndexStore {
    fn merge_segment_data(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        // Merge the term directories
        // The term directory keys are ordered to be most convenient for retrieving all the segments
        // of for a term/field combination in one go (field/term/segment). So we don't end up pulling
        // in a lot of unwanted data, we firstly iterate the keys, it they one of the source segments
        // looking for then we load them and append them to our new segment.

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        let mut current_td_key: Option<(u32, u32)> = None;
        let mut current_td = Vec::new();

        for k in self.db.keys_iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to merge
                break;
            }

            let (field, term, segment) = parse_term_directory_key(&k);

            if !source_segments.contains(&segment) {
                continue;
            }

            if current_td_key != Some((field, term)) {
                // Finished current term directory. Write it to the DB and start the next one
                if let Some((field, term)) = current_td_key {
                    let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
                    if let Err(e) = self.db.put(&kb.key(), &current_td) {
                        return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                    }
                    current_td.clear();
                }

                current_td_key = Some((field, term));
            }

            // Merge term directory into the new one (and remap the doc ids)
            match self.db.get(&k) {
                Ok(Some(docid_set)) => {
                    for doc_id in DocIdSet::FromRDB(docid_set).iter() {
                        let doc_ref = DocRef::from_segment_ord(segment, doc_id);
                        let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();
                        current_td.write_u16::<BigEndian>(*new_doc_id).unwrap();
                    }
                }
                Ok(None) => {},  // FIXME
                Err(e) => return Err(RocksDBReadError::new(k.to_vec(), e).into())
            }
        }

        // All done, write the last term directory
        if let Some((field, term)) = current_td_key {
            let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
            if let Err(e) = self.db.put(&kb.key(), &current_td) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
            current_td.clear();
        }

        // Merge the stored values
        // All stored value keys start with the segment id. So we need to:
        // - Iterate all stored value keys that are prefixed by one of the stored segment ids
        // - Remap their doc ids to the one in the new segment
        // - Write the value back with the new segment/doc ids in the key

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to move
                    break;
                }

                let (segment, doc_id, field, value_type) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                // Remap doc id
                let doc_ref = DocRef::from_segment_ord(segment, doc_id as u16);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();

                // Write value into new segment
                let kb = KeyBuilder::stored_field_value(dest_segment, *new_doc_id, field, &value_type);
                if let Err(e) = self.db.put(&kb.key(), &v) {
                    return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                }
            }
        }

        // Merge the statistics
        // Like stored values, these start with segment ids. But instead of just rewriting the
        // key, we need to sum up all the statistics across the segments being merged.

        let mut statistics = HashMap::new();

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        // Fetch and merge statistics
        for k in self.db.keys_iterator(IteratorMode::From(b"s", Direction::Forward)) {
            if k[0] != b's' {
                // No more statistics to merge
                break;
            }

            let (segment, statistic_name) = parse_statistic_key(&k);

            if !source_segments.contains(&segment) {
                continue;
            }

            match self.db.get(&k) {
                Ok(Some(val_bytes)) => {
                    let value = BigEndian::read_i64(&val_bytes);

                    let mut stat = statistics.entry(statistic_name).or_insert(0);
                    *stat += value;
                }
                Ok(None) => {},  // FIXME
                Err(e) => return Err(RocksDBReadError::new(k.to_vec(), e).into())
            }
        }

        // Write merged statistics to new segment
        for (stat_name, stat_value) in statistics {
            let kb = KeyBuilder::segment_stat(dest_segment, &stat_name);
            let mut val_bytes = [0; 8];
            BigEndian::write_i64(&mut val_bytes, stat_value);
            if let Err(e) = self.db.put(&kb.key(), &val_bytes) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
        }

        // Note: Don't merge the deletion lists
        // Deletion lists can change at any time so we must lock the "document index"
        // before merging them so they can't be altered during merge. we cannot lock
        // this until the commit phase though.

        Ok(())
    }

    fn commit_segment_merge(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        let write_batch = WriteBatch::default();

        // Activate new segment
        let kb = KeyBuilder::segment_active(dest_segment);
        if let Err(e) = write_batch.put(&kb.key(), b"") {
            return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
        }

        // Deactivate old segments
        for source_segment in source_segments.iter() {
            // Activate new segment
            let kb = KeyBuilder::segment_active(*source_segment);
            if let Err(e) = write_batch.delete(&kb.key()) {
                return Err(RocksDBWriteError::new_delete(kb.key().to_vec(), e).into());
            }
        }

        // Update document index and commit
        // This will write the write batch
        try!(self.document_index.commit_segment_merge(&self.db, write_batch, source_segments, dest_segment, doc_ref_mapping));

        Ok(())
    }

    pub fn merge_segments(&self, source_segments: Vec<u32>) -> Result<u32, SegmentMergeError> {
        let dest_segment = try!(self.segments.new_segment(&self.db));

        // Generate a mapping between the ids of the documents in the old segments to the new one
        // This packs the id spaces of the old segments together:
        // For example, say we have to merge 3 segments with 100 documents each:
        //  - The first segment's ids will be the same as before
        //  - The second segment's ids will be remapped to 100 - 199
        //  - The third segment's ids will be remapped to 200 - 299

        let mut doc_ref_mapping: HashMap<DocRef, u16> = HashMap::new();
        let mut current_ord: u32 = 0;

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat(*source_segment, b"total_docs");
            let total_docs = match self.db.get(&kb.key()) {
                Ok(Some(total_docs_bytes)) => {
                    BigEndian::read_i64(&total_docs_bytes)
                }
                Ok(None) => continue,
                Err(e) => return Err(RocksDBReadError::new(kb.key().to_vec(), e).into())
            };

            for source_ord in 0..total_docs {
                if current_ord >= 65536 {
                    return Err(SegmentMergeError::TooManyDocs);
                }

                let from = DocRef::from_segment_ord(*source_segment, source_ord as u16);
                doc_ref_mapping.insert(from, current_ord as u16);
                current_ord += 1;
            }
        }

        // Merge segment data
        // Most of the heavy lifting happens here. This merges all the immutable parts of
        // the segment (which is everything but the deletion list). It does not activate the
        // segment.
        // This means that nothing bad will happen if it crashes half way through -- the
        // worst that could happen is we're left with a partially-written segment that we
        // have to clean up.
        try!(self.merge_segment_data(&source_segments, dest_segment, &doc_ref_mapping));

        // Commit the merge
        // This activates the new segment and updates the document index. Effectively committing
        // the merge.
        // Throughout this stage we need an exclusive lock to the document index. This is to
        // prevent documents in the source segments being deleted/updated so we don't accidentally
        // undelete them (this will block until the merge is complete so they delete/update from
        // the new segment).
        try!(self.commit_segment_merge(&source_segments, dest_segment, &doc_ref_mapping));

        Ok(dest_segment)
    }
}
<commit_msg>Use prefix seek when merging statistics<commit_after>use std::str;
use std::collections::HashMap;

use rocksdb::{Writable, IteratorMode, Direction, WriteBatch};
use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use RocksDBIndexStore;
use errors::{RocksDBReadError, RocksDBWriteError};
use document_index::DocRef;
use key_builder::KeyBuilder;
use search::doc_id_set::DocIdSet;

#[derive(Debug)]
pub enum SegmentMergeError {
    TooManyDocs,
    RocksDBReadError(RocksDBReadError),
    RocksDBWriteError(RocksDBWriteError),
}

impl From<RocksDBReadError> for SegmentMergeError {
    fn from(e: RocksDBReadError) -> SegmentMergeError {
        SegmentMergeError::RocksDBReadError(e)
    }
}

impl From<RocksDBWriteError> for SegmentMergeError {
    fn from(e: RocksDBWriteError) -> SegmentMergeError {
        SegmentMergeError::RocksDBWriteError(e)
    }
}

impl From<SegmentMergeError> for String {
    fn from(e: SegmentMergeError) -> String {
        format!("{:?}", e).to_string()
    }
}

impl RocksDBIndexStore {
    fn merge_segment_data(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        // Merge the term directories
        // The term directory keys are ordered to be most convenient for retrieving all the segments
        // of for a term/field combination in one go (field/term/segment). So we don't end up pulling
        // in a lot of unwanted data, we firstly iterate the keys, it they one of the source segments
        // looking for then we load them and append them to our new segment.

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        let mut current_td_key: Option<(u32, u32)> = None;
        let mut current_td = Vec::new();

        for k in self.db.keys_iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to merge
                break;
            }

            let (field, term, segment) = parse_term_directory_key(&k);

            if !source_segments.contains(&segment) {
                continue;
            }

            if current_td_key != Some((field, term)) {
                // Finished current term directory. Write it to the DB and start the next one
                if let Some((field, term)) = current_td_key {
                    let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
                    if let Err(e) = self.db.put(&kb.key(), &current_td) {
                        return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                    }
                    current_td.clear();
                }

                current_td_key = Some((field, term));
            }

            // Merge term directory into the new one (and remap the doc ids)
            match self.db.get(&k) {
                Ok(Some(docid_set)) => {
                    for doc_id in DocIdSet::FromRDB(docid_set).iter() {
                        let doc_ref = DocRef::from_segment_ord(segment, doc_id);
                        let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();
                        current_td.write_u16::<BigEndian>(*new_doc_id).unwrap();
                    }
                }
                Ok(None) => {},  // FIXME
                Err(e) => return Err(RocksDBReadError::new(k.to_vec(), e).into())
            }
        }

        // All done, write the last term directory
        if let Some((field, term)) = current_td_key {
            let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
            if let Err(e) = self.db.put(&kb.key(), &current_td) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
            current_td.clear();
        }

        // Merge the stored values
        // All stored value keys start with the segment id. So we need to:
        // - Iterate all stored value keys that are prefixed by one of the stored segment ids
        // - Remap their doc ids to the one in the new segment
        // - Write the value back with the new segment/doc ids in the key

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to move
                    break;
                }

                let (segment, doc_id, field, value_type) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                // Remap doc id
                let doc_ref = DocRef::from_segment_ord(segment, doc_id as u16);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();

                // Write value into new segment
                let kb = KeyBuilder::stored_field_value(dest_segment, *new_doc_id, field, &value_type);
                if let Err(e) = self.db.put(&kb.key(), &v) {
                    return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
                }
            }
        }

        // Merge the statistics
        // Like stored values, these start with segment ids. But instead of just rewriting the
        // key, we need to sum up all the statistics across the segments being merged.

        let mut statistics = HashMap::new();

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        // Fetch and merge statistics
        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b's' {
                    // No more statistics to merge
                    break;
                }

                let (segment, statistic_name) = parse_statistic_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                let mut stat = statistics.entry(statistic_name).or_insert(0);
                *stat += BigEndian::read_i64(&v);
            }
        }

        // Write merged statistics to new segment
        for (stat_name, stat_value) in statistics {
            let kb = KeyBuilder::segment_stat(dest_segment, &stat_name);
            let mut val_bytes = [0; 8];
            BigEndian::write_i64(&mut val_bytes, stat_value);
            if let Err(e) = self.db.put(&kb.key(), &val_bytes) {
                return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
            }
        }

        // Note: Don't merge the deletion lists
        // Deletion lists can change at any time so we must lock the "document index"
        // before merging them so they can't be altered during merge. we cannot lock
        // this until the commit phase though.

        Ok(())
    }

    fn commit_segment_merge(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        let write_batch = WriteBatch::default();

        // Activate new segment
        let kb = KeyBuilder::segment_active(dest_segment);
        if let Err(e) = write_batch.put(&kb.key(), b"") {
            return Err(RocksDBWriteError::new_put(kb.key().to_vec(), e).into());
        }

        // Deactivate old segments
        for source_segment in source_segments.iter() {
            // Activate new segment
            let kb = KeyBuilder::segment_active(*source_segment);
            if let Err(e) = write_batch.delete(&kb.key()) {
                return Err(RocksDBWriteError::new_delete(kb.key().to_vec(), e).into());
            }
        }

        // Update document index and commit
        // This will write the write batch
        try!(self.document_index.commit_segment_merge(&self.db, write_batch, source_segments, dest_segment, doc_ref_mapping));

        Ok(())
    }

    pub fn merge_segments(&self, source_segments: Vec<u32>) -> Result<u32, SegmentMergeError> {
        let dest_segment = try!(self.segments.new_segment(&self.db));

        // Generate a mapping between the ids of the documents in the old segments to the new one
        // This packs the id spaces of the old segments together:
        // For example, say we have to merge 3 segments with 100 documents each:
        //  - The first segment's ids will be the same as before
        //  - The second segment's ids will be remapped to 100 - 199
        //  - The third segment's ids will be remapped to 200 - 299

        let mut doc_ref_mapping: HashMap<DocRef, u16> = HashMap::new();
        let mut current_ord: u32 = 0;

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat(*source_segment, b"total_docs");
            let total_docs = match self.db.get(&kb.key()) {
                Ok(Some(total_docs_bytes)) => {
                    BigEndian::read_i64(&total_docs_bytes)
                }
                Ok(None) => continue,
                Err(e) => return Err(RocksDBReadError::new(kb.key().to_vec(), e).into())
            };

            for source_ord in 0..total_docs {
                if current_ord >= 65536 {
                    return Err(SegmentMergeError::TooManyDocs);
                }

                let from = DocRef::from_segment_ord(*source_segment, source_ord as u16);
                doc_ref_mapping.insert(from, current_ord as u16);
                current_ord += 1;
            }
        }

        // Merge segment data
        // Most of the heavy lifting happens here. This merges all the immutable parts of
        // the segment (which is everything but the deletion list). It does not activate the
        // segment.
        // This means that nothing bad will happen if it crashes half way through -- the
        // worst that could happen is we're left with a partially-written segment that we
        // have to clean up.
        try!(self.merge_segment_data(&source_segments, dest_segment, &doc_ref_mapping));

        // Commit the merge
        // This activates the new segment and updates the document index. Effectively committing
        // the merge.
        // Throughout this stage we need an exclusive lock to the document index. This is to
        // prevent documents in the source segments being deleted/updated so we don't accidentally
        // undelete them (this will block until the merge is complete so they delete/update from
        // the new segment).
        try!(self.commit_segment_merge(&source_segments, dest_segment, &doc_ref_mapping));

        Ok(dest_segment)
    }
}
<|endoftext|>


<commit_before>use {osc_types as ot,
    errors as oe};

use std::{io, string, mem, error};
use std::io::BufRead;

use byteorder;
use byteorder::{BigEndian, ReadBytesExt};

/// Common MTP size for ethernet
pub const MTP: usize = 1536;

pub fn decode(msg: &[u8], size: usize) -> ot::OscResult<ot::OscPacket> {
    match msg[0] as char {
        '/' => {
            decode_message(msg, size)
        }
        '#' => {
            decode_bundle(msg)
        }
        _ => Err(oe::OscError::BadOscPacket("Unknown message format.".to_string())),
    }
}

fn decode_message(msg: &[u8], size: usize) -> ot::OscResult<ot::OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);
    let mut pos: u64 = 0;

    match read_osc_string(&mut cursor) {
        Ok(s) => {
            let addr: String = s;
            println!("{}",try!(read_osc_string(&mut cursor)));
        }
        Err(e) => {
            return Err(e)
        }
        Err(e) => return Err(e)
    }

    Ok(ot::OscPacket::Message(ot::OscMessage))
}

fn read_osc_string(cursor: &mut io::Cursor<&[u8]>) -> ot::OscResult<String> {
    let mut str_buf: Vec<u8> = Vec::new();
    match cursor.read_until(0, &mut str_buf) {
        Ok(_) => {
            pad_cursor(cursor);
            String::from_utf8(str_buf)
                .map_err(|e| oe::OscError::StringError(e))
                .map(|s| s.trim_matches(0u8 as char).to_string())
        },
        Err(e) => Err(oe::OscError::ReadError(e)),
    }
}

fn decode_bundle(msg: &[u8]) -> ot::OscResult<ot::OscPacket> {
    Err(oe::OscError::BadOscBundle)
}

fn pad_cursor(cursor: &mut io::Cursor<&[u8]>) {
    let pos = cursor.position();
    match pos % 4 {
        0 => (),
        d => cursor.set_position(pos + (4 - d)),
    }
}<commit_msg>Add reader functions for OSC arguments<commit_after>use {osc_types as ot,
    errors as oe};

use std::{io, string, mem, error};
use std::io::{Read, BufRead};

use byteorder;
use byteorder::{BigEndian, ReadBytesExt};

/// Common MTP size for ethernet
pub const MTP: usize = 1536;

pub fn decode(msg: &[u8], size: usize) -> ot::OscResult<ot::OscPacket> {
    match msg[0] as char {
        '/' => {
            decode_message(msg, size)
        }
        '#' => {
            decode_bundle(msg)
        }
        _ => Err(oe::OscError::BadOscPacket("Unknown message format.".to_string())),
    }
}

fn decode_message(msg: &[u8], size: usize) -> ot::OscResult<ot::OscPacket> {
    let mut cursor: io::Cursor<&[u8]> = io::Cursor::new(msg);
    let mut pos: u64 = 0;

    match read_osc_string(&mut cursor) {
        Ok(s) => {
            let addr: String = s;
            println!("{}",try!(read_osc_string(&mut cursor)));
        }
        Err(e) => {
            return Err(e)
        }
        Err(e) => return Err(e)
    }

    Ok(ot::OscPacket::Message(ot::OscMessage))
}

fn read_osc_string(cursor: &mut io::Cursor<&[u8]>) -> ot::OscResult<String> {
    let mut str_buf: Vec<u8> = Vec::new();
    match cursor.read_until(0, &mut str_buf) {
        Ok(_) => {
            pad_cursor(cursor);
            String::from_utf8(str_buf)
                .map_err(|e| oe::OscError::StringError(e))
                .map(|s| s.trim_matches(0u8 as char).to_string())
        },
        Err(e) => Err(oe::OscError::ReadError(e)),
    }
}

fn read_osc_args(cursor: &mut io::Cursor<&[u8]>, raw_type_tags: String) -> ot::OscResult<Vec<ot::OscType>> {
    let type_tags: Vec<char> = raw_type_tags.chars()
        .skip(1)
        .map(|c| c as char)
        .collect();
    let mut args: Vec<ot::OscType> = Vec::with_capacity(type_tags.len());
    for tag in type_tags {
        match read_osc_arg(cursor, tag) {
            Ok(arg) => {
                args.push(arg);
            },
            Err(e) => return Err(e)
        }
    }
    Ok(args)
}

fn read_osc_arg(cursor: &mut io::Cursor<&[u8]>, tag: char) -> ot::OscResult<ot::OscType> {
    match tag {
        'f' => {
            cursor.read_f32::<BigEndian>()
            .map(|f| ot::OscType::OscFloat(f))
            .map_err(|e| oe::OscError::ByteOrderError(e))
        },
        'd' => {
            cursor.read_f64::<BigEndian>()
            .map(|d| ot::OscType::OscDouble(d))
            .map_err(|e| oe::OscError::ByteOrderError(e))
        }
        'i' => {
            cursor.read_i32::<BigEndian>()
            .map(|i| ot::OscType::OscInt(i))
            .map_err(|e| oe::OscError::ByteOrderError(e))
        },
        'h' => {
            cursor.read_i64::<BigEndian>()
            .map(|l| ot::OscType::OscLong(l))
            .map_err(|e| oe::OscError::ByteOrderError(e))
        },
        's' => {
            read_osc_string(cursor)
            .map(|s| ot::OscType::OscString(s))
        },
        'b' => {
            match cursor.read_u32::<BigEndian>() {
                Ok(size) => {
                    let mut buf: Vec<u8> = Vec::with_capacity(size as usize);
                    match cursor.take(size as u64).read(&mut buf) {
                        Ok(blob_size) => Ok(ot::OscType::OscBlob(buf)),
                        Err(e) => Err(oe::OscError::ReadError(e))
                    }
                },
                Err(e) => return Err(oe::OscError::BadOscBundle)
            }
            // int32 byte count followed by as many bytes
            // read blob ...
        },
        _ => Err(oe::OscError::BadOscArg(format!("Type tag \"{}\" is not implemented!", tag))),
    }
}

fn decode_bundle(msg: &[u8]) -> ot::OscResult<ot::OscPacket> {
    Err(oe::OscError::BadOscBundle)
}

fn pad_cursor(cursor: &mut io::Cursor<&[u8]>) {
    let pos = cursor.position();
    match pos % 4 {
        0 => (),
        d => cursor.set_position(pos + (4 - d)),
    }
}<|endoftext|>


<commit_before>use std::fs;
use std::fs::OpenOptions;
use std::io::prelude::*;
use std::path::PathBuf;
use std::vec::Vec;
use std::str::FromStr;

use time;

#[cfg(test)] use rand;
#[cfg(test)] use rand::Rng;
#[cfg(test)] use test_library::PathPrefix;

use crypto::sha1::Sha1;
use crypto::digest::Digest;

use rustc_serialize::base64::FromBase64;

use message::PackageId;

// TODO: refactor into submodules

pub struct Transfer {
    pub package: PackageId,
    pub checksum: String,
    pub transferred_chunks: Vec<u64>,
    pub prefix_dir: String,
    pub last_chunk_received: i64
}

impl Transfer {
    pub fn new(prefix: String, package: PackageId, checksum: String)
        -> Transfer {
        Transfer {
            package: package,
            checksum: checksum,
            transferred_chunks: Vec::new(),
            prefix_dir: prefix,
            last_chunk_received: time::get_time().sec
        }
    }

    #[cfg(test)]
    pub fn new_test(prefix: &PathPrefix) -> Transfer {
        Transfer {
            package: PackageId {
                name: "".to_string(),
                version: "".to_string()
            },
            checksum: "".to_string(),
            transferred_chunks: Vec::new(),
            prefix_dir: prefix.to_string(),
            last_chunk_received: time::get_time().sec
        }
    }

    #[cfg(test)]
    pub fn randomize(&mut self, i: usize) -> PackageId {
        let name = rand::thread_rng()
            .gen_ascii_chars().take(i).collect::<String>();
        let version = rand::thread_rng()
            .gen_ascii_chars().take(i).collect::<String>();

        trace!("Testing with:");
        trace!("  name: {}\n  version {}", name, version);

        self.package.name = name.clone();
        self.package.version = version.clone();

        PackageId {
            name: name,
            version: version
        }
    }

    pub fn assemble_package(&self) -> bool {
        trace!("Finalizing package {}", self.package);
        try_or!(self.assemble_chunks(), return false);
        self.checksum()
    }

    pub fn write_chunk(&mut self,
                       msg: &str,
                       index: u64) -> bool {
        let success = msg.from_base64().map_err(|e| {
            error!("Could not decode chunk {} for package {}", index, self.package);
            error!("{}", e)
        }).and_then(|msg| self.get_chunk_path(index).map_err(|e| {
            error!("Could not get path for chunk {}", index);
            error!("{}", e)
        }).map(|path| {
            trace!("Saving chunk to {}", path.display());
            if write_new_file(&path, &msg) {
                self.transferred_chunks.push(index);
                self.transferred_chunks.sort();
                self.transferred_chunks.dedup();
                true
            } else {
                error!("Couldn't write chunk {} for package {}", index, self.package);
                false
            }
        })).unwrap_or(false);

        self.last_chunk_received = time::get_time().sec;
        success
    }

    fn get_chunk_path(&self, index: u64) -> Result<PathBuf, String> {
        let mut path = try!(self.get_chunk_dir());
        let filename = index.to_string();

        trace!("Using filename {}", filename);
        path.push(filename);
        Ok(path)
    }

    fn get_package_path(&self) -> Result<PathBuf, String> {
        let mut path = try!(self.get_package_dir());
        path.push(format!("{}.spkg", self.package));
        Ok(path)
    }

    fn assemble_chunks(&self) -> Result<(), String> {
        let package_path = try!(self.get_package_path());

        trace!("Saving package {} to {}", self.package, package_path.display());

        let mut file = try!(OpenOptions::new()
                               .write(true).append(true)
                               .create(true).truncate(true)
                               .open(package_path)
                               .map_err(|x| format!("Couldn't open file: {}", x)));

        let path: PathBuf = try!(self.get_chunk_dir());

        let mut indices = Vec::new();
        for entry in try!(read_dir(&path)) {
            let entry = try!(entry.map_err(|x| format!("No entries: {}", x)));
            let name  = entry.file_name().into_string()
                .unwrap_or("unknown".to_string());

            let chunk_index =
                try!(u64::from_str(&name)
                     .map_err(|_| "Couldn't parse chunk index from filename"
                              .to_string()));
            indices.push(chunk_index);
        }
        indices.sort();

        for index in indices {
            let name = index.to_string();
            let mut chunk_path = path.clone();
            chunk_path.push(&name);
            let mut chunk =
                try!(OpenOptions::new().open(chunk_path)
                     .map_err(|x| format!("Couldn't open file: {}", x)));

            let mut buf = Vec::new();
            try!(chunk.read_to_end(&mut buf)
                 .map_err(|x| format!("Couldn't read file {}: {}", name, x)));
            try!(file.write(&mut buf)
                 .map_err(|x| format!("Couldn't write chunk {} to file {}: {}",
                                      name, self.package, x)));

            trace!("Wrote chunk {} to package {}", name, self.package);
        }
        Ok(())
    }

    fn get_chunk_dir(&self) -> Result<PathBuf, String> {
        let mut path = PathBuf::from(&self.prefix_dir);
        path.push("downloads");
        path.push(format!("{}", self.package));

        fs::create_dir_all(&path).map_err(|e| {
            let path_str = path.to_str().unwrap_or("unknown");
            format!("Couldn't create chunk dir at '{}': {}", path_str, e)
        }).map(|_| path)
    }

    fn checksum(&self) -> bool {
        let path = try_or!(self.get_package_path(), return false);
        let mut file = try_or!(OpenOptions::new().open(path), return false);
        let mut data = Vec::new();

        // TODO: avoid reading in the whole file at once
        try_msg_or!(file.read_to_end(&mut data),
                    "Couldn't read file to check",
                    return false);

        let mut hasher = Sha1::new();
        hasher.input(&data);
        let hash = hasher.result_str();

        if hash == self.checksum {
            true
        } else {
            error!("Checksums didn't match for package {}", self.package);
            error!("    Expected: {}", self.checksum);
            error!("    Got: {}", hash);
            false
        }
    }

    fn get_package_dir(&self) -> Result<PathBuf, String> {
        let mut path = PathBuf::from(&self.prefix_dir);
        path.push("packages");

        fs::create_dir_all(&path).map_err(|e| {
            let path_str = path.to_str().unwrap_or("unknown");
            format!("Couldn't create packges dir at '{}': {}", path_str, e)
        }).map(|_| path)
    }
}

impl Drop for Transfer {
    fn drop(&mut self) {
        let dir = try_or!(self.get_chunk_dir(), return);
        trace!("Dropping transfer for package {}", self.package);

        for entry in try_or!(read_dir(&dir), return) {
            let entry = try_or!(entry, continue);
            let _ = entry.file_name().into_string().map_err(|_|
                error!("Found a malformed entry!")
            ).map(|name| {
                trace!("Dropping chunk file {}", name);
                try_or!(fs::remove_file(entry.path()), return);
            });
        }

        try_or!(fs::remove_dir(dir), return);
    }
}

fn write_new_file(path: &PathBuf, data: &Vec<u8>) -> bool {
    let mut file = try_or!(OpenOptions::new()
                           .write(true).create(true)
                           .truncate(true).open(path),
                           return false);

    try_or!(file.write_all(data), return false);
    try_or!(file.flush(), return false);
    true
}

fn read_dir(path: &PathBuf) -> Result<fs::ReadDir, String> {
    fs::read_dir(path).map_err(|e| {
        let path_str = path.to_str().unwrap_or("unknown");
        format!("Couldn't read dir at '{}': {}", path_str, e)
    })
}

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

    use std::path::PathBuf;
    use std::fs;
    use std::fs::OpenOptions;
    use std::io::prelude::*;

    use rand;
    use rand::Rng;
    use rustc_serialize::base64;
    use rustc_serialize::base64::ToBase64;

    fn create_tmp_directories(prefix: &PathPrefix) {
        for i in 1..20 {
            let mut transfer = Transfer::new_test(prefix);
            let package = transfer.randomize(i);
            let chunk_dir: PathBuf = transfer.get_chunk_dir().unwrap();
            let path = format!("{}/downloads/{}-{}", prefix,
                               package.name, package.version);
            assert_eq!(chunk_dir.to_str().unwrap(), path);

            let path = PathBuf::from(path);
            // This also makes sure it's a directory
            let dir = fs::read_dir(&path).unwrap();

            for _ in dir {
                panic!("Found non-empty directory!");
            }
        }
    }

    #[test]
    fn it_creates_a_tmp_directory() {
        test_init!();
        let prefix = PathPrefix::new();
        create_tmp_directories(&prefix);
    }

    #[test]
    fn it_cleans_up_the_tmp_directories() {
        test_init!();
        let prefix = PathPrefix::new();
        create_tmp_directories(&prefix);
        let path = PathBuf::from(format!("{}/downloads/", prefix));
        let dir = fs::read_dir(&path).unwrap();

        for _ in dir {
            panic!("Found non-empty directory!");
        }
    }

    #[test]
    fn it_creates_a_persistent_directory_per_package() {
        test_init!();
        let prefix = PathPrefix::new();
        for i in 1..20 {
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(i);

            let chunk_dir: PathBuf = transfer.get_package_path().unwrap();
            let path = format!("{}/packages/{}-{}.spkg", prefix,
                               package.name, package.version);
            assert_eq!(chunk_dir.to_str().unwrap(), path);
        }
    }

    macro_rules! assert_chunk_written {
        ($transfer:ident,
         $prefix:ident,
         $package:ident,
         $index:ident,
         $data:ident) => {{
            trace!("Testing with: {}", $data);

            let b64_data = $data.as_bytes().to_base64(
                base64::Config {
                    char_set: base64::CharacterSet::UrlSafe,
                    newline: base64::Newline::LF,
                    pad: true,
                    line_length: None
                });

            trace!("Encoded as: {}", b64_data);

            $transfer.write_chunk(&b64_data, $index as u64);

            let path = format!("{}/downloads/{}-{}/{}", $prefix,
                                $package.name, $package.version, $index);

            trace!("Expecting file at: {}", path);

            let mut from_disk = Vec::new();
            OpenOptions::new()
                .open(PathBuf::from(path))
                .unwrap()
                .read_to_end(&mut from_disk)
                .unwrap();

            assert_eq!($data.into_bytes(), from_disk);
        }}
    }

    #[test]
    fn it_writes_decoded_data_to_disk() {
        test_init!();
        let prefix = PathPrefix::new();
        for i in 1..20 {
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(i);
            for i in 1..20 {
                let data = rand::thread_rng()
                    .gen_ascii_chars().take(i).collect::<String>();
                assert_chunk_written!(transfer, prefix, package, i, data);
            }
        }
    }

    #[test]
    fn it_correctly_assembles_stored_chunks() {
        test_init!();
        let prefix = PathPrefix::new();
        for i in 1..20 {
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(i);
            let mut full_data = String::new();
            for i in 1..20 {
                let data = rand::thread_rng()
                    .gen_ascii_chars().take(i).collect::<String>();
                full_data.push_str(&data);

                assert_chunk_written!(transfer, prefix, package, i, data);
            }

            transfer.assemble_chunks().unwrap();

            let path = format!("{}/packages/{}-{}.spkg", prefix,
                               package.name, package.version);

            trace!("Expecting assembled file at: {}", path);

            let mut from_disk = Vec::new();
            OpenOptions::new()
                .open(PathBuf::from(path))
                .unwrap()
                .read_to_end(&mut from_disk)
                .unwrap();

            assert_eq!(full_data.into_bytes(), from_disk);
        }
    }

    fn checksum_matching(data: String, checksum: String) -> bool {
            let prefix = PathPrefix::new();
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(20);
            let index = 0;
            assert_chunk_written!(transfer, prefix, package, index, data);
            transfer.assemble_chunks().unwrap();

            transfer.checksum = checksum;
            transfer.checksum()
    }

    #[test]
    fn it_returns_true_for_correct_checksums() {
        test_init!();
        assert!(checksum_matching("test\n".to_string(),
        "4e1243bd22c66e76c2ba9eddc1f91394e57f9f83".to_string()));
    }

    #[test]
    fn it_returns_false_for_incorrect_checksums() {
        test_init!();
        assert!(!checksum_matching("test\n".to_string(),
        "fa7c4d75bae3a641d1f9ab5df028175bfb8a69ca".to_string()));
    }

    #[test]
    fn it_returns_false_for_invalid_checksums() {
        test_init!();
        assert!(!checksum_matching("test\n".to_string(),
        "invalid".to_string()));
    }
}
<commit_msg>Refactor persistence module<commit_after>use std::fs;
use std::fs::{OpenOptions, DirEntry, File};
use std::io::prelude::*;
use std::path::PathBuf;
use std::vec::Vec;
use std::str::FromStr;

use time;

#[cfg(test)] use rand;
#[cfg(test)] use rand::Rng;
#[cfg(test)] use test_library::PathPrefix;

use crypto::sha1::Sha1;
use crypto::digest::Digest;

use rustc_serialize::base64::FromBase64;

use message::PackageId;

pub struct Transfer {
    pub package: PackageId,
    pub checksum: String,
    pub transferred_chunks: Vec<u64>,
    pub prefix_dir: String,
    pub last_chunk_received: i64
}

impl Transfer {
    pub fn new(prefix: String, package: PackageId, checksum: String)
        -> Transfer {
        Transfer {
            package: package,
            checksum: checksum,
            transferred_chunks: Vec::new(),
            prefix_dir: prefix,
            last_chunk_received: time::get_time().sec
        }
    }

    #[cfg(test)]
    pub fn new_test(prefix: &PathPrefix) -> Transfer {
        Transfer {
            package: PackageId {
                name: "".to_string(),
                version: "".to_string()
            },
            checksum: "".to_string(),
            transferred_chunks: Vec::new(),
            prefix_dir: prefix.to_string(),
            last_chunk_received: time::get_time().sec
        }
    }

    #[cfg(test)]
    pub fn randomize(&mut self, i: usize) -> PackageId {
        let name = rand::thread_rng()
            .gen_ascii_chars().take(i).collect::<String>();
        let version = rand::thread_rng()
            .gen_ascii_chars().take(i).collect::<String>();

        trace!("Testing with:");
        trace!("  name: {}\n  version {}", name, version);

        self.package.name = name.clone();
        self.package.version = version.clone();

        PackageId {
            name: name,
            version: version
        }
    }

    pub fn write_chunk(&mut self,
                       msg: &str,
                       index: u64) -> bool {
        let success = msg.from_base64().map_err(|e| {
            error!("Could not decode chunk {} for package {}", index, self.package);
            error!("{}", e)
        }).and_then(|msg| self.get_chunk_path(index).map_err(|e| {
            error!("Could not get path for chunk {}", index);
            error!("{}", e)
        }).map(|path| {
            trace!("Saving chunk to {}", path.display());
            if write_new_file(&path, &msg) {
                self.transferred_chunks.push(index);
                self.transferred_chunks.sort();
                self.transferred_chunks.dedup();
                true
            } else {
                error!("Couldn't write chunk {} for package {}", index, self.package);
                false
            }
        })).unwrap_or(false);

        self.last_chunk_received = time::get_time().sec;
        success
    }

    pub fn assemble_package(&self) -> bool {
        trace!("Finalizing package {}", self.package);
        try_or!(self.assemble_chunks(), return false);
        self.checksum()
    }

    fn assemble_chunks(&self) -> Result<(), String> {
        let package_path = try!(self.get_package_path());

        trace!("Saving package {} to {}", self.package, package_path.display());

        let mut file = try!(OpenOptions::new()
                               .write(true).append(true)
                               .create(true).truncate(true)
                               .open(package_path)
                               .map_err(|x| format!("Couldn't open file: {}", x)));

        let path: PathBuf = try!(self.get_chunk_dir());

        // Make sure all indices are valid and sort them
        let mut indices = Vec::new();
        for entry in try!(read_dir(&path)) {
            let entry = try!(entry.map_err(|x| format!("No entries: {}", x)));
            indices.push(try!(parse_index(entry)));
        }
        indices.sort();

        // Append indices to the final file
        for index in indices {
            try!(self.copy_chunk(&path, index, &mut file));
        }
        Ok(())
    }

    /// Reads the chunk file at `path/index` and writes it to `file`
    fn copy_chunk(&self, path: &PathBuf, index: u64, file: &mut File)
        -> Result<(), String> {
        let name = index.to_string();
        let mut chunk_path = path.clone();
        chunk_path.push(&name);
        let mut chunk =
            try!(OpenOptions::new().open(chunk_path)
                 .map_err(|x| format!("Couldn't open file: {}", x)));

        let mut buf = Vec::new();
        try!(chunk.read_to_end(&mut buf)
             .map_err(|x| format!("Couldn't read file {}: {}", name, x)));
        try!(file.write(&mut buf)
             .map_err(|x| format!("Couldn't write chunk {} to file {}: {}",
                                  name, self.package, x)));

        trace!("Wrote chunk {} to package {}", name, self.package);
        Ok(())
    }

    fn checksum(&self) -> bool {
        let path = try_or!(self.get_package_path(), return false);
        let mut file = try_or!(OpenOptions::new().open(path), return false);
        let mut data = Vec::new();

        // TODO: avoid reading in the whole file at once
        try_msg_or!(file.read_to_end(&mut data),
                    "Couldn't read file to check",
                    return false);

        let mut hasher = Sha1::new();
        hasher.input(&data);
        let hash = hasher.result_str();

        if hash == self.checksum {
            true
        } else {
            error!("Checksums didn't match for package {}", self.package);
            error!("    Expected: {}", self.checksum);
            error!("    Got: {}", hash);
            false
        }
    }

    fn get_chunk_path(&self, index: u64) -> Result<PathBuf, String> {
        let mut path = try!(self.get_chunk_dir());
        let filename = index.to_string();

        trace!("Using filename {}", filename);
        path.push(filename);
        Ok(path)
    }

    fn get_package_path(&self) -> Result<PathBuf, String> {
        let mut path = try!(self.get_package_dir());
        path.push(format!("{}.spkg", self.package));
        Ok(path)
    }

    fn get_chunk_dir(&self) -> Result<PathBuf, String> {
        let mut path = PathBuf::from(&self.prefix_dir);
        path.push("downloads");
        path.push(format!("{}", self.package));

        fs::create_dir_all(&path).map_err(|e| {
            let path_str = path.to_str().unwrap_or("unknown");
            format!("Couldn't create chunk dir at '{}': {}", path_str, e)
        }).map(|_| path)
    }

    fn get_package_dir(&self) -> Result<PathBuf, String> {
        let mut path = PathBuf::from(&self.prefix_dir);
        path.push("packages");

        fs::create_dir_all(&path).map_err(|e| {
            let path_str = path.to_str().unwrap_or("unknown");
            format!("Couldn't create packges dir at '{}': {}", path_str, e)
        }).map(|_| path)
    }
}

impl Drop for Transfer {
    fn drop(&mut self) {
        let dir = try_or!(self.get_chunk_dir(), return);
        trace!("Dropping transfer for package {}", self.package);

        for entry in try_or!(read_dir(&dir), return) {
            let entry = try_or!(entry, continue);
            let _ = entry.file_name().into_string().map_err(|_|
                error!("Found a malformed entry!")
            ).map(|name| {
                trace!("Dropping chunk file {}", name);
                try_or!(fs::remove_file(entry.path()), return);
            });
        }

        try_or!(fs::remove_dir(dir), return);
    }
}

fn write_new_file(path: &PathBuf, data: &Vec<u8>) -> bool {
    let mut file = try_or!(OpenOptions::new()
                           .write(true).create(true)
                           .truncate(true).open(path),
                           return false);

    try_or!(file.write_all(data), return false);
    try_or!(file.flush(), return false);
    true
}

fn read_dir(path: &PathBuf) -> Result<fs::ReadDir, String> {
    fs::read_dir(path).map_err(|e| {
        let path_str = path.to_str().unwrap_or("unknown");
        format!("Couldn't read dir at '{}': {}", path_str, e)
    })
}

fn parse_index(entry: DirEntry) -> Result<u64, String> {
    let name = entry.file_name().into_string()
        .unwrap_or("unknown".to_string());
    u64::from_str(&name)
        .map_err(|_| "Couldn't parse chunk index from filename".to_string())
}

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

    use std::path::PathBuf;
    use std::fs;
    use std::fs::OpenOptions;
    use std::io::prelude::*;

    use rand;
    use rand::Rng;
    use rustc_serialize::base64;
    use rustc_serialize::base64::ToBase64;

    fn create_tmp_directories(prefix: &PathPrefix) {
        for i in 1..20 {
            let mut transfer = Transfer::new_test(prefix);
            let package = transfer.randomize(i);
            let chunk_dir: PathBuf = transfer.get_chunk_dir().unwrap();
            let path = format!("{}/downloads/{}-{}", prefix,
                               package.name, package.version);
            assert_eq!(chunk_dir.to_str().unwrap(), path);

            let path = PathBuf::from(path);
            // This also makes sure it's a directory
            let dir = fs::read_dir(&path).unwrap();

            for _ in dir {
                panic!("Found non-empty directory!");
            }
        }
    }

    #[test]
    fn it_creates_a_tmp_directory() {
        test_init!();
        let prefix = PathPrefix::new();
        create_tmp_directories(&prefix);
    }

    #[test]
    fn it_cleans_up_the_tmp_directories() {
        test_init!();
        let prefix = PathPrefix::new();
        create_tmp_directories(&prefix);
        let path = PathBuf::from(format!("{}/downloads/", prefix));
        let dir = fs::read_dir(&path).unwrap();

        for _ in dir {
            panic!("Found non-empty directory!");
        }
    }

    #[test]
    fn it_creates_a_persistent_directory_per_package() {
        test_init!();
        let prefix = PathPrefix::new();
        for i in 1..20 {
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(i);

            let chunk_dir: PathBuf = transfer.get_package_path().unwrap();
            let path = format!("{}/packages/{}-{}.spkg", prefix,
                               package.name, package.version);
            assert_eq!(chunk_dir.to_str().unwrap(), path);
        }
    }

    macro_rules! assert_chunk_written {
        ($transfer:ident,
         $prefix:ident,
         $package:ident,
         $index:ident,
         $data:ident) => {{
            trace!("Testing with: {}", $data);

            let b64_data = $data.as_bytes().to_base64(
                base64::Config {
                    char_set: base64::CharacterSet::UrlSafe,
                    newline: base64::Newline::LF,
                    pad: true,
                    line_length: None
                });

            trace!("Encoded as: {}", b64_data);

            $transfer.write_chunk(&b64_data, $index as u64);

            let path = format!("{}/downloads/{}-{}/{}", $prefix,
                                $package.name, $package.version, $index);

            trace!("Expecting file at: {}", path);

            let mut from_disk = Vec::new();
            OpenOptions::new()
                .open(PathBuf::from(path))
                .unwrap()
                .read_to_end(&mut from_disk)
                .unwrap();

            assert_eq!($data.into_bytes(), from_disk);
        }}
    }

    #[test]
    fn it_writes_decoded_data_to_disk() {
        test_init!();
        let prefix = PathPrefix::new();
        for i in 1..20 {
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(i);
            for i in 1..20 {
                let data = rand::thread_rng()
                    .gen_ascii_chars().take(i).collect::<String>();
                assert_chunk_written!(transfer, prefix, package, i, data);
            }
        }
    }

    #[test]
    fn it_correctly_assembles_stored_chunks() {
        test_init!();
        let prefix = PathPrefix::new();
        for i in 1..20 {
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(i);
            let mut full_data = String::new();
            for i in 1..20 {
                let data = rand::thread_rng()
                    .gen_ascii_chars().take(i).collect::<String>();
                full_data.push_str(&data);

                assert_chunk_written!(transfer, prefix, package, i, data);
            }

            transfer.assemble_chunks().unwrap();

            let path = format!("{}/packages/{}-{}.spkg", prefix,
                               package.name, package.version);

            trace!("Expecting assembled file at: {}", path);

            let mut from_disk = Vec::new();
            OpenOptions::new()
                .open(PathBuf::from(path))
                .unwrap()
                .read_to_end(&mut from_disk)
                .unwrap();

            assert_eq!(full_data.into_bytes(), from_disk);
        }
    }

    fn checksum_matching(data: String, checksum: String) -> bool {
            let prefix = PathPrefix::new();
            let mut transfer = Transfer::new_test(&prefix);
            let package = transfer.randomize(20);
            let index = 0;
            assert_chunk_written!(transfer, prefix, package, index, data);
            transfer.assemble_chunks().unwrap();

            transfer.checksum = checksum;
            transfer.checksum()
    }

    #[test]
    fn it_returns_true_for_correct_checksums() {
        test_init!();
        assert!(checksum_matching("test\n".to_string(),
        "4e1243bd22c66e76c2ba9eddc1f91394e57f9f83".to_string()));
    }

    #[test]
    fn it_returns_false_for_incorrect_checksums() {
        test_init!();
        assert!(!checksum_matching("test\n".to_string(),
        "fa7c4d75bae3a641d1f9ab5df028175bfb8a69ca".to_string()));
    }

    #[test]
    fn it_returns_false_for_invalid_checksums() {
        test_init!();
        assert!(!checksum_matching("test\n".to_string(),
        "invalid".to_string()));
    }
}
<|endoftext|>


<commit_before>#![allow(
    clippy::decimal_literal_representation,
    clippy::derive_partial_eq_without_eq,
    clippy::unreadable_literal,
    clippy::shadow_unrelated
)]

use indoc::indoc;
use serde_derive::{Deserialize, Serialize};
use serde_yaml::{Mapping, Number, Value};
use std::collections::BTreeMap;
use std::fmt::Debug;

fn test_serde<T>(thing: &T, yaml: &str)
where
    T: serde::Serialize + serde::de::DeserializeOwned + PartialEq + Debug,
{
    let serialized = serde_yaml::to_string(&thing).unwrap();
    assert_eq!(yaml, serialized);

    let value = serde_yaml::to_value(&thing).unwrap();
    let serialized = serde_yaml::to_string(&value).unwrap();
    assert_eq!(yaml, serialized);

    let deserialized: T = serde_yaml::from_str(yaml).unwrap();
    assert_eq!(*thing, deserialized);

    let value: Value = serde_yaml::from_str(yaml).unwrap();
    let deserialized = T::deserialize(&value).unwrap();
    assert_eq!(*thing, deserialized);

    let deserialized: T = serde_yaml::from_value(value).unwrap();
    assert_eq!(*thing, deserialized);

    serde_yaml::from_str::<serde::de::IgnoredAny>(yaml).unwrap();
}

#[test]
fn test_default() {
    assert_eq!(Value::default(), Value::Null);
}

#[test]
fn test_int() {
    let thing = 256;
    let yaml = indoc! {"
        256
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_int_max_u64() {
    let thing = u64::MAX;
    let yaml = indoc! {"
        18446744073709551615
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_int_min_i64() {
    let thing = i64::MIN;
    let yaml = indoc! {"
        -9223372036854775808
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_int_max_i64() {
    let thing = i64::MAX;
    let yaml = indoc! {"
        9223372036854775807
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_i128_small() {
    let thing: i128 = -256;
    let yaml = indoc! {"
        -256
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_u128_small() {
    let thing: u128 = 256;
    let yaml = indoc! {"
        256
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_float() {
    let thing = 25.6;
    let yaml = indoc! {"
        25.6
    "};
    test_serde(&thing, yaml);

    let thing = 25.;
    let yaml = indoc! {"
        25.0
    "};
    test_serde(&thing, yaml);

    let thing = f64::INFINITY;
    let yaml = indoc! {"
        .inf
    "};
    test_serde(&thing, yaml);

    let thing = f64::NEG_INFINITY;
    let yaml = indoc! {"
        -.inf
    "};
    test_serde(&thing, yaml);

    let float: f64 = serde_yaml::from_str(indoc! {"
        .nan
    "})
    .unwrap();
    assert!(float.is_nan());
}

#[test]
fn test_float32() {
    let thing: f32 = 25.5;
    let yaml = indoc! {"
        25.5
    "};
    test_serde(&thing, yaml);

    let thing = f32::INFINITY;
    let yaml = indoc! {"
        .inf
    "};
    test_serde(&thing, yaml);

    let thing = f32::NEG_INFINITY;
    let yaml = indoc! {"
        -.inf
    "};
    test_serde(&thing, yaml);

    let single_float: f32 = serde_yaml::from_str(indoc! {"
        .nan
    "})
    .unwrap();
    assert!(single_float.is_nan());
}

#[test]
fn test_char() {
    let ch = '.';
    let yaml = indoc! {"
        '.'
    "};
    assert_eq!(yaml, serde_yaml::to_string(&ch).unwrap());

    let ch = '#';
    let yaml = indoc! {"
        '#'
    "};
    assert_eq!(yaml, serde_yaml::to_string(&ch).unwrap());

    let ch = '-';
    let yaml = indoc! {"
        '-'
    "};
    assert_eq!(yaml, serde_yaml::to_string(&ch).unwrap());
}

#[test]
fn test_vec() {
    let thing = vec![1, 2, 3];
    let yaml = indoc! {"
        - 1
        - 2
        - 3
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_map() {
    let mut thing = BTreeMap::new();
    thing.insert("x".to_owned(), 1);
    thing.insert("y".to_owned(), 2);
    let yaml = indoc! {"
        x: 1
        y: 2
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_basic_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Basic {
        x: isize,
        y: String,
        z: bool,
    }
    let thing = Basic {
        x: -4,
        y: "hi\tquoted".to_owned(),
        z: true,
    };
    let yaml = indoc! {r#"
        x: -4
        y: "hi\tquoted"
        z: true
    "#};
    test_serde(&thing, yaml);
}

#[test]
fn test_multiline_string() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Struct {
        trailing_newline: String,
        no_trailing_newline: String,
    }
    let thing = Struct {
        trailing_newline: "aaa\nbbb\n".to_owned(),
        no_trailing_newline: "aaa\nbbb".to_owned(),
    };
    let yaml = indoc! {r#"
        trailing_newline: |
          aaa
          bbb
        no_trailing_newline: |-
          aaa
          bbb
    "#};
    test_serde(&thing, yaml);
}

#[test]
fn test_strings_needing_quote() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Struct {
        boolean: String,
        integer: String,
        void: String,
    }
    let thing = Struct {
        boolean: "true".to_owned(),
        integer: "1".to_owned(),
        void: "null".to_owned(),
    };
    let yaml = indoc! {r#"
        boolean: 'true'
        integer: '1'
        void: 'null'
    "#};
    test_serde(&thing, yaml);
}

#[test]
fn test_nested_vec() {
    let thing = vec![vec![1, 2, 3], vec![4, 5, 6]];
    let yaml = indoc! {"
        - - 1
          - 2
          - 3
        - - 4
          - 5
          - 6
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_nested_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Outer {
        inner: Inner,
    }
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Inner {
        v: u16,
    }
    let thing = Outer {
        inner: Inner { v: 512 },
    };
    let yaml = indoc! {"
        inner:
          v: 512
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_nested_enum() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Outer {
        Inner(Inner),
    }
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Inner {
        Unit,
    }
    let thing = Outer::Inner(Inner::Unit);
    let yaml = indoc! {"
        !Inner Unit
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_option() {
    let thing = vec![Some(1), None, Some(3)];
    let yaml = indoc! {"
        - 1
        - null
        - 3
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_unit() {
    let thing = vec![(), ()];
    let yaml = indoc! {"
        - null
        - null
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_unit_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Foo;
    let thing = Foo;
    let yaml = indoc! {"
        null
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_unit_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        First,
        Second,
    }
    let thing = Variant::First;
    let yaml = indoc! {"
        First
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_newtype_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct OriginalType {
        v: u16,
    }
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct NewType(OriginalType);
    let thing = NewType(OriginalType { v: 1 });
    let yaml = indoc! {"
        v: 1
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_newtype_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        Size(usize),
    }
    let thing = Variant::Size(127);
    let yaml = indoc! {"
        !Size 127
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_tuple_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        Rgb(u8, u8, u8),
    }
    let thing = Variant::Rgb(32, 64, 96);
    let yaml = indoc! {"
        !Rgb
        - 32
        - 64
        - 96
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_struct_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        Color { r: u8, g: u8, b: u8 },
    }
    let thing = Variant::Color {
        r: 32,
        g: 64,
        b: 96,
    };
    let yaml = indoc! {"
        !Color
        r: 32
        g: 64
        b: 96
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_value() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    pub struct GenericInstructions {
        #[serde(rename = "type")]
        pub typ: String,
        pub config: Value,
    }
    let thing = GenericInstructions {
        typ: "primary".to_string(),
        config: Value::Sequence(vec![
            Value::Null,
            Value::Bool(true),
            Value::Number(Number::from(65535)),
            Value::Number(Number::from(0.54321)),
            Value::String("s".into()),
            Value::Mapping(Mapping::new()),
        ]),
    };
    let yaml = indoc! {"
        type: primary
        config:
        - null
        - true
        - 65535
        - 0.54321
        - s
        - {}
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_mapping() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Data {
        pub substructure: Mapping,
    }

    let mut thing = Data {
        substructure: Mapping::new(),
    };
    thing.substructure.insert(
        Value::String("a".to_owned()),
        Value::String("foo".to_owned()),
    );
    thing.substructure.insert(
        Value::String("b".to_owned()),
        Value::String("bar".to_owned()),
    );

    let yaml = indoc! {"
        substructure:
          a: foo
          b: bar
    "};

    test_serde(&thing, yaml);
}
<commit_msg>Add test of serialize_key/serialize_value map<commit_after>#![allow(
    clippy::decimal_literal_representation,
    clippy::derive_partial_eq_without_eq,
    clippy::unreadable_literal,
    clippy::shadow_unrelated
)]

use indoc::indoc;
use serde::ser::SerializeMap;
use serde_derive::{Deserialize, Serialize};
use serde_yaml::{Mapping, Number, Value};
use std::collections::BTreeMap;
use std::fmt::Debug;

fn test_serde<T>(thing: &T, yaml: &str)
where
    T: serde::Serialize + serde::de::DeserializeOwned + PartialEq + Debug,
{
    let serialized = serde_yaml::to_string(&thing).unwrap();
    assert_eq!(yaml, serialized);

    let value = serde_yaml::to_value(&thing).unwrap();
    let serialized = serde_yaml::to_string(&value).unwrap();
    assert_eq!(yaml, serialized);

    let deserialized: T = serde_yaml::from_str(yaml).unwrap();
    assert_eq!(*thing, deserialized);

    let value: Value = serde_yaml::from_str(yaml).unwrap();
    let deserialized = T::deserialize(&value).unwrap();
    assert_eq!(*thing, deserialized);

    let deserialized: T = serde_yaml::from_value(value).unwrap();
    assert_eq!(*thing, deserialized);

    serde_yaml::from_str::<serde::de::IgnoredAny>(yaml).unwrap();
}

#[test]
fn test_default() {
    assert_eq!(Value::default(), Value::Null);
}

#[test]
fn test_int() {
    let thing = 256;
    let yaml = indoc! {"
        256
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_int_max_u64() {
    let thing = u64::MAX;
    let yaml = indoc! {"
        18446744073709551615
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_int_min_i64() {
    let thing = i64::MIN;
    let yaml = indoc! {"
        -9223372036854775808
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_int_max_i64() {
    let thing = i64::MAX;
    let yaml = indoc! {"
        9223372036854775807
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_i128_small() {
    let thing: i128 = -256;
    let yaml = indoc! {"
        -256
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_u128_small() {
    let thing: u128 = 256;
    let yaml = indoc! {"
        256
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_float() {
    let thing = 25.6;
    let yaml = indoc! {"
        25.6
    "};
    test_serde(&thing, yaml);

    let thing = 25.;
    let yaml = indoc! {"
        25.0
    "};
    test_serde(&thing, yaml);

    let thing = f64::INFINITY;
    let yaml = indoc! {"
        .inf
    "};
    test_serde(&thing, yaml);

    let thing = f64::NEG_INFINITY;
    let yaml = indoc! {"
        -.inf
    "};
    test_serde(&thing, yaml);

    let float: f64 = serde_yaml::from_str(indoc! {"
        .nan
    "})
    .unwrap();
    assert!(float.is_nan());
}

#[test]
fn test_float32() {
    let thing: f32 = 25.5;
    let yaml = indoc! {"
        25.5
    "};
    test_serde(&thing, yaml);

    let thing = f32::INFINITY;
    let yaml = indoc! {"
        .inf
    "};
    test_serde(&thing, yaml);

    let thing = f32::NEG_INFINITY;
    let yaml = indoc! {"
        -.inf
    "};
    test_serde(&thing, yaml);

    let single_float: f32 = serde_yaml::from_str(indoc! {"
        .nan
    "})
    .unwrap();
    assert!(single_float.is_nan());
}

#[test]
fn test_char() {
    let ch = '.';
    let yaml = indoc! {"
        '.'
    "};
    assert_eq!(yaml, serde_yaml::to_string(&ch).unwrap());

    let ch = '#';
    let yaml = indoc! {"
        '#'
    "};
    assert_eq!(yaml, serde_yaml::to_string(&ch).unwrap());

    let ch = '-';
    let yaml = indoc! {"
        '-'
    "};
    assert_eq!(yaml, serde_yaml::to_string(&ch).unwrap());
}

#[test]
fn test_vec() {
    let thing = vec![1, 2, 3];
    let yaml = indoc! {"
        - 1
        - 2
        - 3
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_map() {
    let mut thing = BTreeMap::new();
    thing.insert("x".to_owned(), 1);
    thing.insert("y".to_owned(), 2);
    let yaml = indoc! {"
        x: 1
        y: 2
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_map_key_value() {
    struct Map;

    impl serde::Serialize for Map {
        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where
            S: serde::Serializer,
        {
            // Test maps which do not serialize using serialize_entry.
            let mut map = serializer.serialize_map(Some(1))?;
            map.serialize_key("k")?;
            map.serialize_value("v")?;
            map.end()
        }
    }

    let yaml = indoc! {"
        k: v
    "};
    assert_eq!(yaml, serde_yaml::to_string(&Map).unwrap());
}

#[test]
fn test_basic_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Basic {
        x: isize,
        y: String,
        z: bool,
    }
    let thing = Basic {
        x: -4,
        y: "hi\tquoted".to_owned(),
        z: true,
    };
    let yaml = indoc! {r#"
        x: -4
        y: "hi\tquoted"
        z: true
    "#};
    test_serde(&thing, yaml);
}

#[test]
fn test_multiline_string() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Struct {
        trailing_newline: String,
        no_trailing_newline: String,
    }
    let thing = Struct {
        trailing_newline: "aaa\nbbb\n".to_owned(),
        no_trailing_newline: "aaa\nbbb".to_owned(),
    };
    let yaml = indoc! {r#"
        trailing_newline: |
          aaa
          bbb
        no_trailing_newline: |-
          aaa
          bbb
    "#};
    test_serde(&thing, yaml);
}

#[test]
fn test_strings_needing_quote() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Struct {
        boolean: String,
        integer: String,
        void: String,
    }
    let thing = Struct {
        boolean: "true".to_owned(),
        integer: "1".to_owned(),
        void: "null".to_owned(),
    };
    let yaml = indoc! {r#"
        boolean: 'true'
        integer: '1'
        void: 'null'
    "#};
    test_serde(&thing, yaml);
}

#[test]
fn test_nested_vec() {
    let thing = vec![vec![1, 2, 3], vec![4, 5, 6]];
    let yaml = indoc! {"
        - - 1
          - 2
          - 3
        - - 4
          - 5
          - 6
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_nested_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Outer {
        inner: Inner,
    }
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Inner {
        v: u16,
    }
    let thing = Outer {
        inner: Inner { v: 512 },
    };
    let yaml = indoc! {"
        inner:
          v: 512
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_nested_enum() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Outer {
        Inner(Inner),
    }
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Inner {
        Unit,
    }
    let thing = Outer::Inner(Inner::Unit);
    let yaml = indoc! {"
        !Inner Unit
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_option() {
    let thing = vec![Some(1), None, Some(3)];
    let yaml = indoc! {"
        - 1
        - null
        - 3
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_unit() {
    let thing = vec![(), ()];
    let yaml = indoc! {"
        - null
        - null
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_unit_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Foo;
    let thing = Foo;
    let yaml = indoc! {"
        null
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_unit_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        First,
        Second,
    }
    let thing = Variant::First;
    let yaml = indoc! {"
        First
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_newtype_struct() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct OriginalType {
        v: u16,
    }
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct NewType(OriginalType);
    let thing = NewType(OriginalType { v: 1 });
    let yaml = indoc! {"
        v: 1
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_newtype_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        Size(usize),
    }
    let thing = Variant::Size(127);
    let yaml = indoc! {"
        !Size 127
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_tuple_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        Rgb(u8, u8, u8),
    }
    let thing = Variant::Rgb(32, 64, 96);
    let yaml = indoc! {"
        !Rgb
        - 32
        - 64
        - 96
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_struct_variant() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    enum Variant {
        Color { r: u8, g: u8, b: u8 },
    }
    let thing = Variant::Color {
        r: 32,
        g: 64,
        b: 96,
    };
    let yaml = indoc! {"
        !Color
        r: 32
        g: 64
        b: 96
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_value() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    pub struct GenericInstructions {
        #[serde(rename = "type")]
        pub typ: String,
        pub config: Value,
    }
    let thing = GenericInstructions {
        typ: "primary".to_string(),
        config: Value::Sequence(vec![
            Value::Null,
            Value::Bool(true),
            Value::Number(Number::from(65535)),
            Value::Number(Number::from(0.54321)),
            Value::String("s".into()),
            Value::Mapping(Mapping::new()),
        ]),
    };
    let yaml = indoc! {"
        type: primary
        config:
        - null
        - true
        - 65535
        - 0.54321
        - s
        - {}
    "};
    test_serde(&thing, yaml);
}

#[test]
fn test_mapping() {
    #[derive(Serialize, Deserialize, PartialEq, Debug)]
    struct Data {
        pub substructure: Mapping,
    }

    let mut thing = Data {
        substructure: Mapping::new(),
    };
    thing.substructure.insert(
        Value::String("a".to_owned()),
        Value::String("foo".to_owned()),
    );
    thing.substructure.insert(
        Value::String("b".to_owned()),
        Value::String("bar".to_owned()),
    );

    let yaml = indoc! {"
        substructure:
          a: foo
          b: bar
    "};

    test_serde(&thing, yaml);
}
<|endoftext|>


<commit_before>use super::helpers::fixtures::get_language;
use tree_sitter_tags::{TagKind, TagsConfiguration, TagsContext};

#[test]
fn test_tags_python() {
    let language = get_language("python");
    let tags_config = TagsConfiguration::new(
        language,
        r#"
        ((function_definition
            name: (identifier) @name
            body: (block . (expression_statement (string) @doc))) @function
         (set! strip @doc "(^['\s]*)|(['\s]*$)"))
        (function_definition
            name: (identifier) @name) @function
        ((class_definition
            name: (identifier) @name
            body: (block . (expression_statement (string) @doc))) @class
         (set! strip @doc "(^['\s]*)|(['\s]*$)"))
        (class_definition
            name: (identifier) @name) @class
        (call
            function: (identifier) @name) @call
        "#,
        "",
    )
    .unwrap();

    let mut tag_context = TagsContext::new();
    let tags = tag_context
        .generate_tags(
            &tags_config,
            br#"
        class Customer:
            '''
            Data about a customer
            '''

            def age(self):
                '''
                Get the customer's age
                '''
                compute_age(self.id)
        }
        "#,
        )
        .collect::<Vec<_>>();

    assert_eq!(
        tags.iter().map(|t| (t.name, t.kind)).collect::<Vec<_>>(),
        &[
            ("Customer", TagKind::Class),
            ("age", TagKind::Function),
            ("compute_age", TagKind::Call),
        ]
    );

    assert_eq!(tags[0].docs.as_ref().unwrap(), "Data about a customer");
    assert_eq!(tags[1].docs.as_ref().unwrap(), "Get the customer's age");
}
<commit_msg>Update tags test to reflect new handling of escapes in queries<commit_after>use super::helpers::fixtures::get_language;
use tree_sitter_tags::{TagKind, TagsConfiguration, TagsContext};

#[test]
fn test_tags_python() {
    let language = get_language("python");
    let tags_config = TagsConfiguration::new(
        language,
        r#"
        ((function_definition
            name: (identifier) @name
            body: (block . (expression_statement (string) @doc))) @function
         (set! strip @doc "(^['\"\\s]*)|(['\"\\s]*$)"))
        (function_definition
            name: (identifier) @name) @function
        ((class_definition
            name: (identifier) @name
            body: (block . (expression_statement (string) @doc))) @class
         (set! strip @doc "(^['\"\\s]*)|(['\"\\s]*$)"))
        (class_definition
            name: (identifier) @name) @class
        (call
            function: (identifier) @name) @call
        "#,
        "",
    )
    .unwrap();

    let mut tag_context = TagsContext::new();
    let tags = tag_context
        .generate_tags(
            &tags_config,
            br#"
        class Customer:
            """
            Data about a customer
            """

            def age(self):
                '''
                Get the customer's age
                '''
                compute_age(self.id)
        }
        "#,
        )
        .collect::<Vec<_>>();

    assert_eq!(
        tags.iter().map(|t| (t.name, t.kind)).collect::<Vec<_>>(),
        &[
            ("Customer", TagKind::Class),
            ("age", TagKind::Function),
            ("compute_age", TagKind::Call),
        ]
    );

    assert_eq!(tags[0].docs.as_ref().unwrap(), "Data about a customer");
    assert_eq!(tags[1].docs.as_ref().unwrap(), "Get the customer's age");
}
<|endoftext|>


<commit_before>extern crate htmlescape;

use std::marker::PhantomData;
use std::time::Duration;

use irc::client::prelude::*;

use regex::Regex;

use plugin::*;
use utils::Url;
use FrippyClient;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::Fail;
use failure::ResultExt;

lazy_static! {
    static ref URL_RE: Regex = Regex::new(r"(^|\s)(https?://\S+)").unwrap();
    static ref WORD_RE: Regex = Regex::new(r"(\w+)").unwrap();
}

#[derive(PluginName, Debug)]
pub struct UrlTitles<C> {
    max_kib: usize,
    phantom: PhantomData<C>,
}

#[derive(Clone, Debug)]
struct Title(String, Option<usize>);

impl From<String> for Title {
    fn from(title: String) -> Self {
        Title(title, None)
    }
}

impl From<Title> for String {
    fn from(title: Title) -> Self {
        title.0
    }
}

impl Title {
    fn find_by_delimiters(body: &str, delimiters: [&str; 3]) -> Result<Self, UrlError> {
        let title = body
            .find(delimiters[0])
            .map(|tag| {
                body[tag..]
                    .find(delimiters[1])
                    .map(|offset| tag + offset + delimiters[1].len())
                    .map(|start| {
                        body[start..]
                            .find(delimiters[2])
                            .map(|offset| start + offset)
                            .map(|end| &body[start..end])
                    })
            })
            .and_then(|s| s.and_then(|s| s))
            .ok_or(ErrorKind::MissingTitle)?;

        debug!("Found title {:?} with delimiters {:?}", title, delimiters);

        htmlescape::decode_html(title)
            .map(|t| t.into())
            .map_err(|_| ErrorKind::HtmlDecoding.into())
    }

    fn find_ogtitle(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["property=\"og:title\"", "content=\"", "\""])
    }

    fn find_title(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["<title", ">", "</title>"])
    }

    // TODO Improve logic
    fn get_usefulness(self, url: &str) -> Self {
        let mut usefulness = 0;
        for word in WORD_RE.find_iter(&self.0) {
            let w = word.as_str().to_lowercase();
            if w.len() > 2 && !url.to_lowercase().contains(&w) {
                usefulness += 1;
            }
        }

        Title(self.0, Some(usefulness))
    }

    pub fn usefulness(&self) -> usize {
        self.1.expect("Usefulness should be calculated already")
    }

    fn clean_up(self) -> Self {
        Title(self.0.trim().replace('\n', "|").replace('\r', "|"), self.1)
    }

    pub fn find_clean_ogtitle(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_ogtitle(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }

    pub fn find_clean_title(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_title(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }
}

impl<C: FrippyClient> UrlTitles<C> {
    /// If a file is larger than `max_kib` KiB the download is stopped
    pub fn new(max_kib: usize) -> Self {
        UrlTitles {
            max_kib,
            phantom: PhantomData,
        }
    }

    fn grep_url<'a>(&self, msg: &'a str) -> Option<Url<'a>> {
        let captures = URL_RE.captures(msg)?;
        debug!("Url captures: {:?}", captures);

        Some(captures.get(2)?.as_str().into())
    }

    fn url(&self, text: &str) -> Result<String, UrlError> {
        let url = self
            .grep_url(text)
            .ok_or(ErrorKind::MissingUrl)?
            .max_kib(self.max_kib)
            .timeout(Duration::from_secs(5));
        let body = url.request().context(ErrorKind::Download)?;

        let title = Title::find_clean_title(&body, url.as_str());
        let og_title = Title::find_clean_ogtitle(&body, url.as_str());

        let title = match (title, og_title) {
            (Ok(title), Ok(og_title)) => {
                if title.usefulness() > og_title.usefulness() {
                    title
                } else {
                    og_title
                }
            }
            (Ok(title), _) => title,
            (_, Ok(title)) => title,
            (Err(e), _) => Err(e)?,
        };

        Ok(title.into())
    }
}

impl<C: FrippyClient> Plugin for UrlTitles<C> {
    type Client = C;
    fn execute(&self, _: &Self::Client, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref msg) => {
                if URL_RE.is_match(msg) {
                    ExecutionStatus::RequiresThread
                } else {
                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(
        &self,
        client: &Self::Client,
        message: &Message,
    ) -> Result<(), FrippyError> {
        if let Command::PRIVMSG(_, ref content) = message.command {
            let title = self.url(content).context(FrippyErrorKind::Url)?;
            let response = format!("[URL] {}", title);

            client
                .send_privmsg(message.response_target().unwrap(), &response)
                .context(FrippyErrorKind::Connection)?;
        }

        Ok(())
    }

    fn command(&self, client: &Self::Client, command: PluginCommand) -> Result<(), FrippyError> {
        client
            .send_notice(
                &command.source,
                "This Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?;

        Ok(())
    }

    fn evaluate(&self, _: &Self::Client, command: PluginCommand) -> Result<String, String> {
        self.url(&command.tokens[0])
            .map_err(|e| e.cause().unwrap().to_string())
    }
}

pub mod error {
    /// A URL plugin error
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "UrlError"]
    pub enum ErrorKind {
        /// A download error
        #[fail(display = "A download error occured")]
        Download,

        /// Missing URL error
        #[fail(display = "No URL was found")]
        MissingUrl,

        /// Missing title error
        #[fail(display = "No title was found")]
        MissingTitle,

        /// Useless title error
        #[fail(display = "The titles found were not useful enough")]
        UselessTitle,

        /// Html decoding error
        #[fail(display = "Failed to decode Html characters")]
        HtmlDecoding,
    }
}
<commit_msg>Url: Fix useless titles getting printed<commit_after>extern crate htmlescape;

use std::marker::PhantomData;
use std::time::Duration;

use irc::client::prelude::*;

use regex::Regex;

use plugin::*;
use utils::Url;
use FrippyClient;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::Fail;
use failure::ResultExt;

lazy_static! {
    static ref URL_RE: Regex = Regex::new(r"(^|\s)(https?://\S+)").unwrap();
    static ref WORD_RE: Regex = Regex::new(r"(\w+)").unwrap();
}

#[derive(PluginName, Debug)]
pub struct UrlTitles<C> {
    max_kib: usize,
    phantom: PhantomData<C>,
}

#[derive(Clone, Debug)]
struct Title(String, Option<usize>);

impl From<String> for Title {
    fn from(title: String) -> Self {
        Title(title, None)
    }
}

impl From<Title> for String {
    fn from(title: Title) -> Self {
        title.0
    }
}

impl Title {
    fn find_by_delimiters(body: &str, delimiters: [&str; 3]) -> Result<Self, UrlError> {
        let title = body
            .find(delimiters[0])
            .map(|tag| {
                body[tag..]
                    .find(delimiters[1])
                    .map(|offset| tag + offset + delimiters[1].len())
                    .map(|start| {
                        body[start..]
                            .find(delimiters[2])
                            .map(|offset| start + offset)
                            .map(|end| &body[start..end])
                    })
            })
            .and_then(|s| s.and_then(|s| s))
            .ok_or(ErrorKind::MissingTitle)?;

        debug!("Found title {:?} with delimiters {:?}", title, delimiters);

        htmlescape::decode_html(title)
            .map(|t| t.into())
            .map_err(|_| ErrorKind::HtmlDecoding.into())
    }

    fn find_ogtitle(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["property=\"og:title\"", "content=\"", "\""])
    }

    fn find_title(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["<title", ">", "</title>"])
    }

    // TODO Improve logic
    fn get_usefulness(self, url: &str) -> Self {
        let mut usefulness = 0;
        for word in WORD_RE.find_iter(&self.0) {
            let w = word.as_str().to_lowercase();
            if w.len() > 2 && !url.to_lowercase().contains(&w) {
                usefulness += 1;
            }
        }

        Title(self.0, Some(usefulness))
    }

    pub fn usefulness(&self) -> usize {
        self.1.expect("Usefulness should be calculated already")
    }

    fn clean_up(self) -> Self {
        Title(self.0.trim().replace('\n', "|").replace('\r', "|"), self.1)
    }

    pub fn find_clean_ogtitle(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_ogtitle(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }

    pub fn find_clean_title(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_title(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }
}

impl<C: FrippyClient> UrlTitles<C> {
    /// If a file is larger than `max_kib` KiB the download is stopped
    pub fn new(max_kib: usize) -> Self {
        UrlTitles {
            max_kib,
            phantom: PhantomData,
        }
    }

    fn grep_url<'a>(&self, msg: &'a str) -> Option<Url<'a>> {
        let captures = URL_RE.captures(msg)?;
        debug!("Url captures: {:?}", captures);

        Some(captures.get(2)?.as_str().into())
    }

    fn url(&self, text: &str) -> Result<String, UrlError> {
        let url = self
            .grep_url(text)
            .ok_or(ErrorKind::MissingUrl)?
            .max_kib(self.max_kib)
            .timeout(Duration::from_secs(5));
        let body = url.request().context(ErrorKind::Download)?;

        let title = Title::find_clean_title(&body, url.as_str());
        let og_title = Title::find_clean_ogtitle(&body, url.as_str());

        let title = match (title, og_title) {
            (Ok(title), Ok(og_title)) => {
                if title.usefulness() > og_title.usefulness() {
                    title
                } else {
                    og_title
                }
            }
            (Ok(title), _) => title,
            (_, Ok(title)) => title,
            (Err(e), _) => Err(e)?,
        };

        if title.usefulness() == 0 {
            Err(ErrorKind::UselessTitle)?;
        }

        Ok(title.into())
    }
}

impl<C: FrippyClient> Plugin for UrlTitles<C> {
    type Client = C;
    fn execute(&self, _: &Self::Client, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref msg) => {
                if URL_RE.is_match(msg) {
                    ExecutionStatus::RequiresThread
                } else {
                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(
        &self,
        client: &Self::Client,
        message: &Message,
    ) -> Result<(), FrippyError> {
        if let Command::PRIVMSG(_, ref content) = message.command {
            let title = self.url(content).context(FrippyErrorKind::Url)?;
            let response = format!("[URL] {}", title);

            client
                .send_privmsg(message.response_target().unwrap(), &response)
                .context(FrippyErrorKind::Connection)?;
        }

        Ok(())
    }

    fn command(&self, client: &Self::Client, command: PluginCommand) -> Result<(), FrippyError> {
        client
            .send_notice(
                &command.source,
                "This Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?;

        Ok(())
    }

    fn evaluate(&self, _: &Self::Client, command: PluginCommand) -> Result<String, String> {
        self.url(&command.tokens[0])
            .map_err(|e| e.cause().unwrap().to_string())
    }
}

pub mod error {
    /// A URL plugin error
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "UrlError"]
    pub enum ErrorKind {
        /// A download error
        #[fail(display = "A download error occured")]
        Download,

        /// Missing URL error
        #[fail(display = "No URL was found")]
        MissingUrl,

        /// Missing title error
        #[fail(display = "No title was found")]
        MissingTitle,

        /// Useless title error
        #[fail(display = "The titles found were not useful enough")]
        UselessTitle,

        /// Html decoding error
        #[fail(display = "Failed to decode Html characters")]
        HtmlDecoding,
    }
}
<|endoftext|>