<commit_before>// Copyright 2016 Pierre Marijon.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! GFF3 format reading and writing.
//!
//! GFF2 definition : http://gmod.org/wiki/GFF2#The_GFF2_File_Format (not yet support)
//! GTF2 definition : http://mblab.wustl.edu/GTF2.html (not yet support)
//! GFF3 definition : http://gmod.org/wiki/GFF3#GFF3_Format
//!
//! # Example
//!
//! ```
//! use std::io;
//! use bio::io::gff;
//! let reader = gff::Reader::new(io::stdin());
//! ```

use std::io;
use std::fs;
use std::path::Path;
use std::convert::AsRef;
use std::collections::HashMap;

use csv;

use io::Strand;

/// A GFF reader.
pub struct Reader<R: io::Read> {
    inner: csv::Reader<R>,
}

impl Reader<fs::File> {
    /// Read GFF from given file path.
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::new)
    }
}

impl<R: io::Read> Reader<R> {
    /// Create a new GFF reader given an instance of `io::Read`.
    pub fn new(reader: R) -> Self {
        Reader {
            inner: csv::Reader::from_reader(reader).delimiter(b'\t').has_headers(false)
        }
    }

    /// Iterate over all records.
    pub fn records(&mut self) -> Records<R> {
        Records { inner: self.inner.decode() }
    }
}

/// A GFF record.
pub struct Records<'a, R: 'a + io::Read> {
    inner: csv::DecodedRecords<'a, R, (String, String, String, u64, u64, String, String, String, String)>,
}

impl<'a, R: io::Read> Iterator for Records<'a, R> {
    type Item = csv::Result<Record>;

    fn next(&mut self) -> Option<csv::Result<Record>> {
        self.inner.next().map(|res| {
            res.map(|(seqname, source, feature_type, start, end, score, strand, frame, attributes)| {
                Record {
                    seqname: seqname,
                    source: source,
                    feature_type: feature_type,
                    start: start,
                    end: end,
                    score: score,
                    strand: strand,
                    frame: frame,
                    attributes: csv::Reader::from_string(attributes)
                        .delimiter(b'=')
                        .record_terminator(csv::RecordTerminator::Any(b';'))
                        .has_headers(false)
                        .decode().collect::<csv::Result<HashMap<String, String>>>().unwrap(),
                }
            })
        })
    }
}

/// A GFF writer.
pub struct Writer<W: io::Write> {
    inner: csv::Writer<W>,
}

impl Writer<fs::File> {
    /// Write to a given file path.
    pub fn to_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::create(path).map(Writer::new)
    }
}

impl<W: io::Write> Writer<W> {
    /// Write to a given writer.
    pub fn new(writer: W) -> Self {
        Writer { inner: csv::Writer::from_writer(writer).delimiter(b'\t').flexible(true) }
    }

    /// Write a given GFF record.
    pub fn write(&mut self, record: Record) -> csv::Result<()> {
        if !record.attributes.is_empty() {
            self.inner.encode((
                record.seqname,
                record.source,
                record.feature_type,
                record.start,
                record.end,
                record.score,
                record.strand,
                record.frame,
                record.attributes.iter().map(|(a, b)| format!("{}={}", a, b)).collect::<Vec<_>>().join(";"),
                ))

        } else {
            self.inner.encode((record.seqname, record.source, record.feature_type, record.start, record.end, record.score, record.strand, record.frame, ""))
        }
    }
}

/// A GFF record
#[derive(RustcEncodable)]
pub struct Record {
    seqname: String,
    source: String,
    feature_type: String,
    start: u64,
    end: u64,
    score: String,
    strand: String,
    frame: String,
    attributes: HashMap<String, String>,
}

impl Record {
    /// Create a new GFF record.
    pub fn new() -> Self {
        Record {
            seqname: "".to_owned(),
            source: "".to_owned(),
            feature_type: "".to_owned(),
            start: 0,
            end: 0,
            score: ".".to_owned(),
            strand: ".".to_owned(),
            frame: "".to_owned(),
            attributes: HashMap::<String, String>::new(),
        }
    }

    /// Sequence name of the feature.
    pub fn seqname(&self) -> &str {
        &self.seqname
    }

    /// Source of the feature.
    pub fn source(&self) -> &str {
        &self.source
    }

    /// Type of the feature.
    pub fn feature_type(&self) -> &str {
        &self.feature_type
    }

    /// Start position of feature (1-based).
    pub fn start(&self) -> u64 {
        self.start
    }

    /// End position of feature (1-based, not included).
    pub fn end(&self) -> u64 {
        self.end
    }

    /// Score of feature
    pub fn score(&self) -> Option<u64> {
        match self.score.as_ref() {
            "." => None,
            _ => self.score.parse::<u64>().ok(),
        }
    }

    /// Strand of the feature.
    pub fn strand(&self) -> Option<Strand> {
        match self.strand.as_ref() {
            "+" => Some(Strand::Forward),
            "-" => Some(Strand::Reverse),
            _ => None,
        }
    }

    /// Frame of the feature.
    pub fn frame(&self) -> &str {
        &self.frame
    }

    /// Attribute of feature
    pub fn attributes(&self) -> &HashMap<String, String> {
        &self.attributes
    }

    /// Set seqname of feature.
    pub fn set_seqname(&mut self, seqname: &str) {
        self.seqname = seqname.to_owned();
    }

    /// Set source of feature.
    pub fn set_source(&mut self, source: &str) {
        self.source = source.to_owned();
    }

    /// Set type of feature.
    pub fn set_feature_type(&mut self, feature_type: &str) {
        self.feature_type = feature_type.to_owned();
    }

    /// Set start of feature.
    pub fn set_start(&mut self, start: u64) {
        self.start = start;
    }

    /// Set end of feature.
    pub fn set_end(&mut self, end: u64) {
        self.end = end;
    }

    /// Set score of feature.
    pub fn set_score(&mut self, score: &str) {
        self.score = score.to_owned();
    }

    /// Set strand of feature.
    pub fn set_strand(&mut self, strand: &str) {
        self.strand = strand.to_owned();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use io::Strand;
    use std::collections::HashMap;

    const GFF_FILE: &'static [u8] = b"P0A7B8\tUniProtKB\tInitiator methionine\t1\t1\t.\t.\t.\tNote=Removed;ID=test
P0A7B8\tUniProtKB\tChain\t2\t176\t50\t+\t.\tNote=ATP-dependent protease subunit HslV;ID=PRO_0000148105
";
    //required because HashMap iter on element randomly
    const GFF_FILE_NO_ATTRIB: &'static [u8] = b"P0A7B8\tUniProtKB\tInitiator methionine\t1\t1\t.\t.\t.\t
P0A7B8\tUniProtKB\tChain\t2\t176\t50\t+\t.\t
";

    #[test]
    fn test_reader() {
        let seqname = ["P0A7B8", "P0A7B8"];
        let source = ["UniProtKB", "UniProtKB"];
        let feature_type = ["Initiator methionine", "Chain"];
        let starts = [1, 2];
        let ends = [1, 176];
        let scores = [None, Some(50)];
        let strand = [None, Some(Strand::Forward)];
        let frame = [".", "."];
        let mut attributes = [HashMap::new(), HashMap::new()];
        attributes[0].insert("ID".to_owned(), "test".to_owned());
        attributes[0].insert("Note".to_owned(), "Removed".to_owned());
        attributes[1].insert("ID".to_owned(), "PRO_0000148105".to_owned());
        attributes[1].insert("Note".to_owned(), "ATP-dependent protease subunit HslV".to_owned());

        let mut reader = Reader::new(GFF_FILE);
        for (i, r) in reader.records().enumerate() {
            let record = r.ok().expect("Error reading record");
            assert_eq!(record.seqname(), seqname[i]);
            assert_eq!(record.source(), source[i]);
            assert_eq!(record.feature_type(), feature_type[i]);
            assert_eq!(record.start(), starts[i]);
            assert_eq!(record.end(), ends[i]);
            assert_eq!(record.score(), scores[i]);
            assert_eq!(record.strand(), strand[i]);
            assert_eq!(record.frame(), frame[i]);
            assert_eq!(record.attributes(), &attributes[i]);
        }
    }

    #[test]
    fn test_writer() {
        let mut reader = Reader::new(GFF_FILE_NO_ATTRIB);
        let mut writer = Writer::new(vec![]);
        for r in reader.records() {
            writer.write(r.ok().expect("Error reading record")).ok().expect("Error writing record");
        }
        assert_eq!(writer.inner.as_string(), String::from_utf8_lossy(GFF_FILE_NO_ATTRIB))
    }
}
<commit_msg>More clear and pretty write function.<commit_after>// Copyright 2016 Pierre Marijon.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! GFF3 format reading and writing.
//!
//! GFF2 definition : http://gmod.org/wiki/GFF2#The_GFF2_File_Format (not yet support)
//! GTF2 definition : http://mblab.wustl.edu/GTF2.html (not yet support)
//! GFF3 definition : http://gmod.org/wiki/GFF3#GFF3_Format
//!
//! # Example
//!
//! ```
//! use std::io;
//! use bio::io::gff;
//! let reader = gff::Reader::new(io::stdin());
//! ```

use std::io;
use std::fs;
use std::path::Path;
use std::convert::AsRef;
use std::collections::HashMap;

use csv;

use io::Strand;

/// A GFF reader.
pub struct Reader<R: io::Read> {
    inner: csv::Reader<R>,
}

impl Reader<fs::File> {
    /// Read GFF from given file path.
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::new)
    }
}

impl<R: io::Read> Reader<R> {
    /// Create a new GFF reader given an instance of `io::Read`.
    pub fn new(reader: R) -> Self {
        Reader {
            inner: csv::Reader::from_reader(reader).delimiter(b'\t').has_headers(false)
        }
    }

    /// Iterate over all records.
    pub fn records(&mut self) -> Records<R> {
        Records { inner: self.inner.decode() }
    }
}

/// A GFF record.
pub struct Records<'a, R: 'a + io::Read> {
    inner: csv::DecodedRecords<'a, R, (String, String, String, u64, u64, String, String, String, String)>,
}

impl<'a, R: io::Read> Iterator for Records<'a, R> {
    type Item = csv::Result<Record>;

    fn next(&mut self) -> Option<csv::Result<Record>> {
        self.inner.next().map(|res| {
            res.map(|(seqname, source, feature_type, start, end, score, strand, frame, attributes)| {
                Record {
                    seqname: seqname,
                    source: source,
                    feature_type: feature_type,
                    start: start,
                    end: end,
                    score: score,
                    strand: strand,
                    frame: frame,
                    attributes: csv::Reader::from_string(attributes)
                        .delimiter(b'=')
                        .record_terminator(csv::RecordTerminator::Any(b';'))
                        .has_headers(false)
                        .decode().collect::<csv::Result<HashMap<String, String>>>().unwrap(),
                }
            })
        })
    }
}

/// A GFF writer.
pub struct Writer<W: io::Write> {
    inner: csv::Writer<W>,
}

impl Writer<fs::File> {
    /// Write to a given file path.
    pub fn to_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::create(path).map(Writer::new)
    }
}

impl<W: io::Write> Writer<W> {
    /// Write to a given writer.
    pub fn new(writer: W) -> Self {
        Writer { inner: csv::Writer::from_writer(writer).delimiter(b'\t').flexible(true) }
    }

    /// Write a given GFF record.
    pub fn write(&mut self, record: Record) -> csv::Result<()> {
        let attributes;
        if !record.attributes.is_empty() {
            attributes = record.attributes.iter().map(|(a, b)| format!("{}={}", a, b)).collect::<Vec<_>>().join(";");
        } else {
            attributes = "".to_owned();
        }
        self.inner.encode((record.seqname, record.source, record.feature_type, record.start, record.end, record.score, record.strand, record.frame, attributes))
    }
}

/// A GFF record
#[derive(RustcEncodable)]
pub struct Record {
    seqname: String,
    source: String,
    feature_type: String,
    start: u64,
    end: u64,
    score: String,
    strand: String,
    frame: String,
    attributes: HashMap<String, String>,
}

impl Record {
    /// Create a new GFF record.
    pub fn new() -> Self {
        Record {
            seqname: "".to_owned(),
            source: "".to_owned(),
            feature_type: "".to_owned(),
            start: 0,
            end: 0,
            score: ".".to_owned(),
            strand: ".".to_owned(),
            frame: "".to_owned(),
            attributes: HashMap::<String, String>::new(),
        }
    }

    /// Sequence name of the feature.
    pub fn seqname(&self) -> &str {
        &self.seqname
    }

    /// Source of the feature.
    pub fn source(&self) -> &str {
        &self.source
    }

    /// Type of the feature.
    pub fn feature_type(&self) -> &str {
        &self.feature_type
    }

    /// Start position of feature (1-based).
    pub fn start(&self) -> u64 {
        self.start
    }

    /// End position of feature (1-based, not included).
    pub fn end(&self) -> u64 {
        self.end
    }

    /// Score of feature
    pub fn score(&self) -> Option<u64> {
        match self.score.as_ref() {
            "." => None,
            _ => self.score.parse::<u64>().ok(),
        }
    }

    /// Strand of the feature.
    pub fn strand(&self) -> Option<Strand> {
        match self.strand.as_ref() {
            "+" => Some(Strand::Forward),
            "-" => Some(Strand::Reverse),
            _ => None,
        }
    }

    /// Frame of the feature.
    pub fn frame(&self) -> &str {
        &self.frame
    }

    /// Attribute of feature
    pub fn attributes(&self) -> &HashMap<String, String> {
        &self.attributes
    }

    /// Set seqname of feature.
    pub fn set_seqname(&mut self, seqname: &str) {
        self.seqname = seqname.to_owned();
    }

    /// Set source of feature.
    pub fn set_source(&mut self, source: &str) {
        self.source = source.to_owned();
    }

    /// Set type of feature.
    pub fn set_feature_type(&mut self, feature_type: &str) {
        self.feature_type = feature_type.to_owned();
    }

    /// Set start of feature.
    pub fn set_start(&mut self, start: u64) {
        self.start = start;
    }

    /// Set end of feature.
    pub fn set_end(&mut self, end: u64) {
        self.end = end;
    }

    /// Set score of feature.
    pub fn set_score(&mut self, score: &str) {
        self.score = score.to_owned();
    }

    /// Set strand of feature.
    pub fn set_strand(&mut self, strand: &str) {
        self.strand = strand.to_owned();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use io::Strand;
    use std::collections::HashMap;

    const GFF_FILE: &'static [u8] = b"P0A7B8\tUniProtKB\tInitiator methionine\t1\t1\t.\t.\t.\tNote=Removed;ID=test
P0A7B8\tUniProtKB\tChain\t2\t176\t50\t+\t.\tNote=ATP-dependent protease subunit HslV;ID=PRO_0000148105
";
    //required because HashMap iter on element randomly
    const GFF_FILE_NO_ATTRIB: &'static [u8] = b"P0A7B8\tUniProtKB\tInitiator methionine\t1\t1\t.\t.\t.\t
P0A7B8\tUniProtKB\tChain\t2\t176\t50\t+\t.\t
";

    #[test]
    fn test_reader() {
        let seqname = ["P0A7B8", "P0A7B8"];
        let source = ["UniProtKB", "UniProtKB"];
        let feature_type = ["Initiator methionine", "Chain"];
        let starts = [1, 2];
        let ends = [1, 176];
        let scores = [None, Some(50)];
        let strand = [None, Some(Strand::Forward)];
        let frame = [".", "."];
        let mut attributes = [HashMap::new(), HashMap::new()];
        attributes[0].insert("ID".to_owned(), "test".to_owned());
        attributes[0].insert("Note".to_owned(), "Removed".to_owned());
        attributes[1].insert("ID".to_owned(), "PRO_0000148105".to_owned());
        attributes[1].insert("Note".to_owned(), "ATP-dependent protease subunit HslV".to_owned());

        let mut reader = Reader::new(GFF_FILE);
        for (i, r) in reader.records().enumerate() {
            let record = r.ok().expect("Error reading record");
            assert_eq!(record.seqname(), seqname[i]);
            assert_eq!(record.source(), source[i]);
            assert_eq!(record.feature_type(), feature_type[i]);
            assert_eq!(record.start(), starts[i]);
            assert_eq!(record.end(), ends[i]);
            assert_eq!(record.score(), scores[i]);
            assert_eq!(record.strand(), strand[i]);
            assert_eq!(record.frame(), frame[i]);
            assert_eq!(record.attributes(), &attributes[i]);
        }
    }

    #[test]
    fn test_writer() {
        let mut reader = Reader::new(GFF_FILE_NO_ATTRIB);
        let mut writer = Writer::new(vec![]);
        for r in reader.records() {
            writer.write(r.ok().expect("Error reading record")).ok().expect("Error writing record");
        }
        assert_eq!(writer.inner.as_string(), String::from_utf8_lossy(GFF_FILE_NO_ATTRIB))
    }
}
<|endoftext|>


<commit_before>use strum_macros::{Display};
use super::arch::{Arch, Arch32, rva_plus_usize};

use failure::{Error, Fail};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
}

#[derive(Display, Clone, Copy)]
pub enum DetectedArch {
    X32,
    X64
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw,  // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

pub struct Section<A: Arch> {
    pub addr: A::RVA,
    pub buf: Vec<u8>,
    // TODO
    pub perms: u8,
    pub name: String,
}

impl <A: Arch> Section<A> {
    pub fn contains(self: &Section<A>, rva: A::RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        if let Some(max) = rva_plus_usize::<A>(self.addr, self.buf.len()) {
            if rva >= max {
                return false
            }
        } else {
            return false
        }

        true
    }
}

pub struct LoadedModule<A: Arch> {
    pub base_address: A::VA,
    pub sections: Vec<Section<A>>,
}

pub trait Loader {
    fn get_arch(&self) -> DetectedArch;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!("{}/{}/{}", self.get_plat(), self.get_arch(), self.get_file_format());
    }

    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given buffer as a 32-bit module.
    ///
    /// Panics if this isn't a loader that supports 32-bit modules.
    fn load32(&self, buf: &[u8]) -> Result<LoadedModule<Arch32>, Error>;
}

// TODO: implement loaders/x32/windows/pe
// TODO: implement loaders/x64/windows/pe

pub struct ShellcodeLoader {
    arch: DetectedArch,
    plat: Platform,
}

impl ShellcodeLoader {
    pub fn new(arch: DetectedArch, plat: Platform) -> ShellcodeLoader {
        ShellcodeLoader {
            arch,
            plat,
        }
    }
}

impl Loader for ShellcodeLoader {
    fn get_arch(&self) -> DetectedArch {
        self.arch
    }

    fn get_plat(&self) -> Platform {
        self.plat
    }

    fn get_file_format(&self) -> FileFormat {
        FileFormat::Raw
    }

    fn taste(&self, _buf: &[u8]) -> bool {
        // we can load anything as shellcode
        true
    }

    /// ```
    /// use lancelot::loader::*;
    ///
    /// let loader = lancelot::loader::ShellcodeLoader::new(DetectedArch::X32, Platform::Windows);
    /// loader.load32(b"MZ\x90\x00")
    ///   .map(|module| {
    ///     assert_eq!(module.base_address,     0x0);
    ///     assert_eq!(module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    fn load32(&self, buf: &[u8]) -> Result<LoadedModule<Arch32>, Error> {
        if let DetectedArch::X32 = self.arch {
            Ok(LoadedModule::<Arch32>{
                base_address: 0x0,
                sections: vec![
                    Section::<Arch32> {
                        addr: 0x0,
                        buf: buf.to_vec(),
                        perms: 0x0, // TODO
                        name: "raw".to_string(),
                    }
                ]
            })
        } else {
            panic!("not a 32-bit loader")
        }
    }
}

pub fn default_loaders() -> Vec<Box<dyn Loader>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader>> = vec![];
    // the order here matters!
    // the default `load32` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.
    loaders.push(Box::new(ShellcodeLoader::new(DetectedArch::X32, Platform::Windows)));
    loaders.push(Box::new(ShellcodeLoader::new(DetectedArch::X64, Platform::Windows)));
    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::loader::*;
///
/// match taste(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/X32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste(buf: &[u8]) -> impl Iterator<Item=Box<dyn Loader>> {
    default_loaders()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample as a 32-bit module using the first matching
///  loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::loader::*;
///
/// load32(b"\xEB\xFE")
///   .map(|module| {
///     assert_eq!(module.base_address,     0x0);
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load32(buf: &[u8]) -> Result<LoadedModule<Arch32>, Error> {
    match taste(buf).nth(0) {
        Some(loader) => {
            loader.load32(buf)
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<commit_msg>loader: when invoking `load32` helper, return tuple with the loader, too<commit_after>use strum_macros::{Display};
use super::arch::{Arch, Arch32, rva_plus_usize};

use failure::{Error, Fail};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
}

#[derive(Display, Clone, Copy)]
pub enum DetectedArch {
    X32,
    X64,
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw,  // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

pub struct Section<A: Arch> {
    pub addr: A::RVA,
    pub buf: Vec<u8>,
    // TODO
    pub perms: u8,
    pub name: String,
}

impl <A: Arch> Section<A> {
    pub fn contains(self: &Section<A>, rva: A::RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        if let Some(max) = rva_plus_usize::<A>(self.addr, self.buf.len()) {
            if rva >= max {
                return false
            }
        } else {
            return false
        }

        true
    }
}

pub struct LoadedModule<A: Arch> {
    pub base_address: A::VA,
    pub sections: Vec<Section<A>>,
}

pub trait Loader {
    fn get_arch(&self) -> DetectedArch;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!("{}/{}/{}", self.get_plat(), self.get_arch(), self.get_file_format());
    }

    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given buffer as a 32-bit module.
    ///
    /// Panics if this isn't a loader that supports 32-bit modules.
    fn load32(&self, buf: &[u8]) -> Result<LoadedModule<Arch32>, Error>;
}

// TODO: implement loaders/x32/windows/pe
// TODO: implement loaders/x64/windows/pe

pub struct ShellcodeLoader {
    arch: DetectedArch,
    plat: Platform,
}

impl ShellcodeLoader {
    pub fn new(arch: DetectedArch, plat: Platform) -> ShellcodeLoader {
        ShellcodeLoader {
            arch,
            plat,
        }
    }
}

impl Loader for ShellcodeLoader {
    fn get_arch(&self) -> DetectedArch {
        self.arch
    }

    fn get_plat(&self) -> Platform {
        self.plat
    }

    fn get_file_format(&self) -> FileFormat {
        FileFormat::Raw
    }

    fn taste(&self, _buf: &[u8]) -> bool {
        // we can load anything as shellcode
        true
    }

    /// ```
    /// use lancelot::loader::*;
    ///
    /// let loader = lancelot::loader::ShellcodeLoader::new(DetectedArch::X32, Platform::Windows);
    /// loader.load32(b"MZ\x90\x00")
    ///   .map(|module| {
    ///     assert_eq!(module.base_address,     0x0);
    ///     assert_eq!(module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    fn load32(&self, buf: &[u8]) -> Result<LoadedModule<Arch32>, Error> {
        if let DetectedArch::X32 = self.arch {
            Ok(LoadedModule::<Arch32>{
                base_address: 0x0,
                sections: vec![
                    Section::<Arch32> {
                        addr: 0x0,
                        buf: buf.to_vec(),
                        perms: 0x0, // TODO
                        name: "raw".to_string(),
                    }
                ]
            })
        } else {
            panic!("not a 32-bit loader")
        }
    }
}

pub fn default_loaders() -> Vec<Box<dyn Loader>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader>> = vec![];
    // the order here matters!
    // the default `load32` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.
    loaders.push(Box::new(ShellcodeLoader::new(DetectedArch::X32, Platform::Windows)));
    loaders.push(Box::new(ShellcodeLoader::new(DetectedArch::X64, Platform::Windows)));
    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::loader::*;
///
/// match taste(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/X32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste(buf: &[u8]) -> impl Iterator<Item=Box<dyn Loader>> {
    default_loaders()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample as a 32-bit module using the first matching
///  loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::loader::*;
///
/// load32(b"\xEB\xFE")
///   .map(|(loader, module)| {
///     assert_eq!(loader.get_name(),       "Windows/X32/Raw");
///     assert_eq!(module.base_address,     0x0);
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load32(buf: &[u8]) -> Result<(Box<dyn Loader>, LoadedModule<Arch32>), Error> {
    match taste(buf).nth(0) {
        Some(loader) => {
            loader.load32(buf).map(|module|
                (loader, module)
            )
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<|endoftext|>


<commit_before>#[macro_export]
macro_rules! compile_prim(
    ($ty:ty, $type_name:ident, $make_constant:ident) => (
impl Compile for $ty {
    #[inline(always)]
    fn compile<'a>(&self, func:&UncompiledFunction<'a>) -> Value<'a> {
        unsafe {
            NativeRef::from_ptr($make_constant(func.as_ptr(), $type_name, *self) )
        }
    }
    #[inline(always)]
    fn jit_type(_:Option<$ty>) -> Type {
        unsafe {
            NativeRef::from_ptr($type_name)
        }
    }
});
    ($ty:ty, $type_name:ident, $make_constant:ident, $cast:ty) => (
impl Compile for $ty {
    #[inline(always)]
    fn compile<'a>(&self, func:&UncompiledFunction<'a>) -> Value<'a> {
        unsafe {
            NativeRef::from_ptr($make_constant(func.as_ptr(), $type_name, *self as $cast) )
        }
    }
    #[inline(always)]
    fn jit_type(_:Option<$ty>) -> Type {
        unsafe {
            NativeRef::from_ptr($type_name)
        }
    }
});
)

macro_rules! compile_func(
    (fn($($arg:ident),*) -> $ret:ty, $sig:ty) => (
    impl<$($arg:Compile,)* R:Compile> Compile for $sig {
        #[inline(always)]
        fn compile<'a>(&self, func:&UncompiledFunction<'a>) -> Value<'a> {
            let ptr = (self as *const $sig).to_uint().compile(func);
            func.insn_convert(&ptr, get::<$sig>(), false)
        }
        #[inline(always)]
        fn jit_type(_:Option<$sig>) -> Type {
            Type::create_signature(CDECL, get::<R>(), [$(get::<$arg>()),*][mut])
        }
    })
)

macro_rules! compile_prims(
    ($(($ty:ty, $cast: ty) => ($type_name:ident, $make_constant:ident)),+) => (
        $(compile_prim!($ty, $type_name, $make_constant, $cast))+
    );
)

#[macro_export]
macro_rules! native_ref(
    ($name:ident, $field:ident, $pointer_ty:ty) => (
        #[deriving(PartialEq)]
        pub struct $name {
            $field: $pointer_ty
        }
        impl NativeRef for $name {
            #[inline(always)]
            /// Convert to a native pointer
            unsafe fn as_ptr(&self) -> $pointer_ty {
                self.$field
            }
            #[inline(always)]
            /// Convert from a native pointer
            unsafe fn from_ptr(ptr:$pointer_ty) -> $name {
                $name {
                    $field: ptr
                }
            }
        }
    );
    ($name:ident, $field:ident, $pointer_ty:ty, $lifetime:ident) => (
        #[deriving(PartialEq)]
        pub struct $name<'a> {
            $field: $pointer_ty,
            marker: $lifetime<'a>
        }
        impl<'a> NativeRef for $name<'a> {
            #[inline]
            /// Convert to a native pointer
            unsafe fn as_ptr(&self) -> $pointer_ty {
                self.$field
            }
            #[inline]
            /// Convert from a native pointer
            unsafe fn from_ptr<'a>(ptr:$pointer_ty) -> $name<'a> {
                $name {
                    $field: ptr,
                    marker: $lifetime::<'a>
                }
            }
        }
    )
)
#[macro_export]
macro_rules! jit(
    (struct {
        $($name:expr: $ty:ty),*
    }) => ({
        let structure = Type::create_struct([
            $(get::<$ty>()),*
        ].as_mut_slice());
        structure.set_names(&[$($name),*]);
        structure
    });
    (union {
        $($name:expr: $ty:ty),+
    }) => ({
        let structure = Type::create_union([
            $(get::<$ty>()),+
        ].as_mut_slice());
        structure.set_names(&[$($name),+]);
        structure
    });
    ($func:ident, return) => (
        $func.insn_default_return()
    );
    ($func:ident, return $($t:tt)+) => (
        $func.insn_return(jit!($($t)+))
    );
    ($func:ident, $var:ident += $num:expr) => (
        $func.insn_store($var, &$func.insn_add($var, $num));
    );
    ($func:ident, $var:ident -= $num:expr) => (
        $func.insn_store($var, &$func.insn_sub($var, $num));
    );
    ($func:ident, $var:ident *= $num:expr) => (
        $func.insn_store($var, &$func.insn_mul($var, $num));
    );
    ($func:ident, $var:ident /= $num:expr) => (
        $func.insn_store($var, &$func.insn_div($var, $num));
    );
    ($func:ident, $a:expr + $b:expr) => (
        $func.insn_add($a, $b)
    );
    ($func:ident, $a:expr - $b:expr) => (
        $func.insn_sub($a, $b)
    );
    ($func:ident, $a:expr * $b:expr) => (
        $func.insn_mul($a, $b)
    );
    ($func:ident, $a:expr / $b:expr) => (
        $func.insn_div($a, $b)
    );
    ($func:ident, $a:expr % $b:expr) => (
        $func.insn_rem($a, $b)
    );
    ($func:ident, $var:ident = $val:expr) => (
        $func.insn_store($var, $val);
    );
    ($func:ident, *$var:ident) => (
        $func.insn_load($var)
    );
    ($func:ident, call($call:expr,
        $($arg:expr),+
    )) => (
        $func.insn_call(None::<String>, $call, None, [$($arg),+].as_mut_slice())
    );
    ($func:ident, jump_table($value:expr,
        $($label:ident),+
    )) => (
    let ($($label),+) = {
        $(let $label:Label = Label::new($func);)+
        $func.insn_jump_table($value, [
            $($label),+
        ].as_mut_slice());
        ($($label),+)
    });
    ($func:expr, $value:expr) => (
        $func.insn_of(&$value)
    );
)<commit_msg>ALlow missing copy in Context<commit_after>#[macro_export]
macro_rules! compile_prim(
    ($ty:ty, $type_name:ident, $make_constant:ident) => (
impl Compile for $ty {
    #[inline(always)]
    fn compile<'a>(&self, func:&UncompiledFunction<'a>) -> Value<'a> {
        unsafe {
            NativeRef::from_ptr($make_constant(func.as_ptr(), $type_name, *self) )
        }
    }
    #[inline(always)]
    fn jit_type(_:Option<$ty>) -> Type {
        unsafe {
            NativeRef::from_ptr($type_name)
        }
    }
});
    ($ty:ty, $type_name:ident, $make_constant:ident, $cast:ty) => (
impl Compile for $ty {
    #[inline(always)]
    fn compile<'a>(&self, func:&UncompiledFunction<'a>) -> Value<'a> {
        unsafe {
            NativeRef::from_ptr($make_constant(func.as_ptr(), $type_name, *self as $cast) )
        }
    }
    #[inline(always)]
    fn jit_type(_:Option<$ty>) -> Type {
        unsafe {
            NativeRef::from_ptr($type_name)
        }
    }
});
)

macro_rules! compile_func(
    (fn($($arg:ident),*) -> $ret:ty, $sig:ty) => (
    impl<$($arg:Compile,)* R:Compile> Compile for $sig {
        #[inline(always)]
        fn compile<'a>(&self, func:&UncompiledFunction<'a>) -> Value<'a> {
            let ptr = (self as *const $sig).to_uint().compile(func);
            func.insn_convert(&ptr, get::<$sig>(), false)
        }
        #[inline(always)]
        fn jit_type(_:Option<$sig>) -> Type {
            Type::create_signature(CDECL, get::<R>(), [$(get::<$arg>()),*][mut])
        }
    })
)

macro_rules! compile_prims(
    ($(($ty:ty, $cast: ty) => ($type_name:ident, $make_constant:ident)),+) => (
        $(compile_prim!($ty, $type_name, $make_constant, $cast))+
    );
)

#[macro_export]
macro_rules! native_ref(
    ($name:ident, $field:ident, $pointer_ty:ty) => (
        #[allow(missing_copy_implementations)]
        #[deriving(PartialEq)]
        pub struct $name {
            $field: $pointer_ty
        }
        impl NativeRef for $name {
            #[inline(always)]
            /// Convert to a native pointer
            unsafe fn as_ptr(&self) -> $pointer_ty {
                self.$field
            }
            #[inline(always)]
            /// Convert from a native pointer
            unsafe fn from_ptr(ptr:$pointer_ty) -> $name {
                $name {
                    $field: ptr
                }
            }
        }
    );
    ($name:ident, $field:ident, $pointer_ty:ty, $lifetime:ident) => (
        #[deriving(PartialEq)]
        pub struct $name<'a> {
            $field: $pointer_ty,
            marker: $lifetime<'a>
        }
        impl<'a> NativeRef for $name<'a> {
            #[inline]
            /// Convert to a native pointer
            unsafe fn as_ptr(&self) -> $pointer_ty {
                self.$field
            }
            #[inline]
            /// Convert from a native pointer
            unsafe fn from_ptr<'a>(ptr:$pointer_ty) -> $name<'a> {
                $name {
                    $field: ptr,
                    marker: $lifetime::<'a>
                }
            }
        }
    )
)
#[macro_export]
macro_rules! jit(
    (struct {
        $($name:expr: $ty:ty),*
    }) => ({
        let structure = Type::create_struct([
            $(get::<$ty>()),*
        ].as_mut_slice());
        structure.set_names(&[$($name),*]);
        structure
    });
    (union {
        $($name:expr: $ty:ty),+
    }) => ({
        let structure = Type::create_union([
            $(get::<$ty>()),+
        ].as_mut_slice());
        structure.set_names(&[$($name),+]);
        structure
    });
    ($func:ident, return) => (
        $func.insn_default_return()
    );
    ($func:ident, return $($t:tt)+) => (
        $func.insn_return(jit!($($t)+))
    );
    ($func:ident, $var:ident += $num:expr) => (
        $func.insn_store($var, &$func.insn_add($var, $num));
    );
    ($func:ident, $var:ident -= $num:expr) => (
        $func.insn_store($var, &$func.insn_sub($var, $num));
    );
    ($func:ident, $var:ident *= $num:expr) => (
        $func.insn_store($var, &$func.insn_mul($var, $num));
    );
    ($func:ident, $var:ident /= $num:expr) => (
        $func.insn_store($var, &$func.insn_div($var, $num));
    );
    ($func:ident, $a:expr + $b:expr) => (
        $func.insn_add($a, $b)
    );
    ($func:ident, $a:expr - $b:expr) => (
        $func.insn_sub($a, $b)
    );
    ($func:ident, $a:expr * $b:expr) => (
        $func.insn_mul($a, $b)
    );
    ($func:ident, $a:expr / $b:expr) => (
        $func.insn_div($a, $b)
    );
    ($func:ident, $a:expr % $b:expr) => (
        $func.insn_rem($a, $b)
    );
    ($func:ident, $var:ident = $val:expr) => (
        $func.insn_store($var, $val);
    );
    ($func:ident, *$var:ident) => (
        $func.insn_load($var)
    );
    ($func:ident, call($call:expr,
        $($arg:expr),+
    )) => (
        $func.insn_call(None::<String>, $call, None, [$($arg),+].as_mut_slice())
    );
    ($func:ident, jump_table($value:expr,
        $($label:ident),+
    )) => (
    let ($($label),+) = {
        $(let $label:Label = Label::new($func);)+
        $func.insn_jump_table($value, [
            $($label),+
        ].as_mut_slice());
        ($($label),+)
    });
    ($func:expr, $value:expr) => (
        $func.insn_of(&$value)
    );
)<|endoftext|>


<commit_before>use common;
use runner;
use serde_yaml;
use std::{io, str};
use std::collections::BTreeMap;
use std::fs::File;

#[derive(Debug, PartialEq, Eq)]
pub enum Token {
    Junk,
    Comma,
    FullStop,
    QuestionMark,
    OpenBracket,
    CloseBracket,
    Colon,
    Quote,
    SingleQuote,
    Newline,
    Word(String),
    Begin,
}

impl Token {
    pub fn string(self: Token) -> String {
        return match self {
                Token::Word(word) => word,
                Token::Junk => String::from(" "),
                Token::Comma => String::from(","),
                Token::FullStop => String::from("."),
                Token::QuestionMark => String::from("?"),
                Token::OpenBracket => String::from("("),
                Token::CloseBracket => String::from(")"),
                Token::Colon => String::from(":"),
                Token::Quote => String::from("\""),
                Token::SingleQuote => String::from("\'"),
                Token::Newline => String::from("\n"),
                Token::Begin => String::from(""),
            }
            .to_lowercase();
    }
}

fn is_alphabetic(c: char) -> bool {
    c.is_alphabetic()
}

fn alpha_or_word_chars(c: char) -> bool {
    c.is_alphabetic() || c == '\'' || c == '’' || c == '-' || c == '\u{2014}' // em-dash
}

named!(word_match<&str, (&str, &str)>,
    tuple!(
        take_while1!(is_alphabetic),
        take_while!(alpha_or_word_chars)
    )
);

named!(get_token<&str, Token>,
    alt!(
        tag_s!(".") => {|_| Token::FullStop} |
        tag_s!(",") => {|_| Token::Comma} |
        tag_s!(":") => {|_| Token::Colon} |
        tag_s!("?") => {|_| Token::QuestionMark} |
        tag_s!("(") => {|_| Token::OpenBracket} |
        tag_s!(")") => {|_| Token::CloseBracket} |
        is_a_s!("\"“”") => {|_| Token::Quote} |
        is_a_s!("\r\n") => {|_| Token::Newline} |
        is_a_s!("\'") => {|_| Token::SingleQuote} |
        word_match => { |(begin, rest)| {
                let mut word = String::from(begin);
                word += rest;
                Token::Word(word)
            }
        } |
        take_s!( 1 )  => { |_| Token::Junk }
    )
);

fn empty_filter(x: Token) -> Option<Token> {
    Some(x)
}

pub fn get_tokens(filename: &str) -> Result<Vec<Token>, io::Error> {
    return common::get_words_core_fn(filename, get_token, empty_filter);
}

#[derive(Serialize, Deserialize, Debug)]
pub struct MarkovInfo {
    scores: BTreeMap<u8, BTreeMap<String, u16>>,
    lookup: BTreeMap<String, BTreeMap<u8, BTreeMap<String, u8>>>,
}

impl MarkovInfo {
    fn new() -> MarkovInfo {
        MarkovInfo {
            scores: BTreeMap::new(),
            lookup: BTreeMap::new(),
        }
    }

    fn add_token(self: &mut MarkovInfo, last: String, token: &String) {
        let last_hash = self.lookup.entry(last.clone()).or_insert(BTreeMap::new());
        let token_score = runner::score(&token);
        let score_entry = last_hash.entry(token_score).or_insert(BTreeMap::new());
        *score_entry.entry(token.clone()).or_insert(0) += 1;
        let root_score_entry = self.scores.entry(token_score).or_insert(BTreeMap::new());
        *root_score_entry.entry(token.clone()).or_insert(0) += 1;
    }

    fn default_get(self: &MarkovInfo, score: u8) -> String {
        match self.scores.get(&score) {
            Some(score_hash) => score_hash.keys().nth(0).unwrap().clone(),
            None => {
                warn!("Have no words with score {} so making up one", score);
                let mut ret = String::new();
                for _ in 0..(score / 10) {
                    ret.push('z');
                }
                for _ in 0..(score % 10) {
                    ret.push('a');
                }
                return ret;
            }
        }
    }

    fn get_token(self: &MarkovInfo, last: &String, score: u8) -> String {
        debug!("Looking up for '{}' and {}", last, score);
        return match self.lookup.get(last) {
            Some(word) => {
                match word.get(&score) {
                    Some(score_hash) => score_hash.keys().nth(0).unwrap().clone(),
                    None => self.default_get(score),
                }
            }
            None => self.default_get(score),
        };
    }
}

pub fn generate_markov<'a>(filename: &str) -> Result<MarkovInfo, io::Error> {
    let tokens = get_tokens(filename)?;
    let str_tokens = tokens.into_iter()
        .filter(|t: &Token| -> bool { if let &Token::Junk = t { false } else { true } })
        .map(|t| t.string());
    let mut res = MarkovInfo::new();
    let mut last = Token::Begin.string();
    for token in str_tokens {
        res.add_token(last, &token);
        last = token.clone();
    }
    println!("{:?}", res.scores.get(&0).unwrap());
    Ok(res)
}

pub fn make_beatnik(wottasquare: &str, markov_fname: &str) -> Result<String, io::Error> {
    let buffer = File::open(markov_fname).unwrap();
    let markov: MarkovInfo = serde_yaml::from_reader(&buffer).map_err(common::io_str_error)?;
    let words = runner::get_wottas(wottasquare)?;
    let mut last = Token::Begin.string();
    let mut out = String::new();
    for word in words {
        let token = markov.get_token(&last, word.score());
        out.push_str(&format!("{} ", token));
        last = token;
    }
    return Ok(out);
}

#[cfg(test)]
mod tests {
    use common;
    use quickcheck::TestResult;
    use super::{empty_filter, get_token};

    #[test]
    fn copes_with_utf_8() {
        common::get_words_core("why—I".as_bytes(), get_token, empty_filter).unwrap();
    }

    quickcheck! {
      fn token_test(xs: String) -> TestResult {
          return match common::get_words_core(xs.as_bytes(), get_token, empty_filter) {
              Ok(_) => TestResult::passed(),
              Err(err) => {
                  println!("Error: '{}'", xs);
                  TestResult::error(format!("{:?}", err))
              }
          }
      }
  }
}
<commit_msg>Don't need to output the zero scores from Markov generation<commit_after>use common;
use runner;
use serde_yaml;
use std::{io, str};
use std::collections::BTreeMap;
use std::fs::File;

#[derive(Debug, PartialEq, Eq)]
pub enum Token {
    Junk,
    Comma,
    FullStop,
    QuestionMark,
    OpenBracket,
    CloseBracket,
    Colon,
    Quote,
    SingleQuote,
    Newline,
    Word(String),
    Begin,
}

impl Token {
    pub fn string(self: Token) -> String {
        return match self {
                Token::Word(word) => word,
                Token::Junk => String::from(" "),
                Token::Comma => String::from(","),
                Token::FullStop => String::from("."),
                Token::QuestionMark => String::from("?"),
                Token::OpenBracket => String::from("("),
                Token::CloseBracket => String::from(")"),
                Token::Colon => String::from(":"),
                Token::Quote => String::from("\""),
                Token::SingleQuote => String::from("\'"),
                Token::Newline => String::from("\n"),
                Token::Begin => String::from(""),
            }
            .to_lowercase();
    }
}

fn is_alphabetic(c: char) -> bool {
    c.is_alphabetic()
}

fn alpha_or_word_chars(c: char) -> bool {
    c.is_alphabetic() || c == '\'' || c == '’' || c == '-' || c == '\u{2014}' // em-dash
}

named!(word_match<&str, (&str, &str)>,
    tuple!(
        take_while1!(is_alphabetic),
        take_while!(alpha_or_word_chars)
    )
);

named!(get_token<&str, Token>,
    alt!(
        tag_s!(".") => {|_| Token::FullStop} |
        tag_s!(",") => {|_| Token::Comma} |
        tag_s!(":") => {|_| Token::Colon} |
        tag_s!("?") => {|_| Token::QuestionMark} |
        tag_s!("(") => {|_| Token::OpenBracket} |
        tag_s!(")") => {|_| Token::CloseBracket} |
        is_a_s!("\"“”") => {|_| Token::Quote} |
        is_a_s!("\r\n") => {|_| Token::Newline} |
        is_a_s!("\'") => {|_| Token::SingleQuote} |
        word_match => { |(begin, rest)| {
                let mut word = String::from(begin);
                word += rest;
                Token::Word(word)
            }
        } |
        take_s!( 1 )  => { |_| Token::Junk }
    )
);

fn empty_filter(x: Token) -> Option<Token> {
    Some(x)
}

pub fn get_tokens(filename: &str) -> Result<Vec<Token>, io::Error> {
    return common::get_words_core_fn(filename, get_token, empty_filter);
}

#[derive(Serialize, Deserialize, Debug)]
pub struct MarkovInfo {
    scores: BTreeMap<u8, BTreeMap<String, u16>>,
    lookup: BTreeMap<String, BTreeMap<u8, BTreeMap<String, u8>>>,
}

impl MarkovInfo {
    fn new() -> MarkovInfo {
        MarkovInfo {
            scores: BTreeMap::new(),
            lookup: BTreeMap::new(),
        }
    }

    fn add_token(self: &mut MarkovInfo, last: String, token: &String) {
        let last_hash = self.lookup.entry(last.clone()).or_insert(BTreeMap::new());
        let token_score = runner::score(&token);
        let score_entry = last_hash.entry(token_score).or_insert(BTreeMap::new());
        *score_entry.entry(token.clone()).or_insert(0) += 1;
        let root_score_entry = self.scores.entry(token_score).or_insert(BTreeMap::new());
        *root_score_entry.entry(token.clone()).or_insert(0) += 1;
    }

    fn default_get(self: &MarkovInfo, score: u8) -> String {
        match self.scores.get(&score) {
            Some(score_hash) => score_hash.keys().nth(0).unwrap().clone(),
            None => {
                warn!("Have no words with score {} so making up one", score);
                let mut ret = String::new();
                for _ in 0..(score / 10) {
                    ret.push('z');
                }
                for _ in 0..(score % 10) {
                    ret.push('a');
                }
                return ret;
            }
        }
    }

    fn get_token(self: &MarkovInfo, last: &String, score: u8) -> String {
        debug!("Looking up for '{}' and {}", last, score);
        return match self.lookup.get(last) {
            Some(word) => {
                match word.get(&score) {
                    Some(score_hash) => score_hash.keys().nth(0).unwrap().clone(),
                    None => self.default_get(score),
                }
            }
            None => self.default_get(score),
        };
    }
}

pub fn generate_markov<'a>(filename: &str) -> Result<MarkovInfo, io::Error> {
    let tokens = get_tokens(filename)?;
    let str_tokens = tokens.into_iter()
        .filter(|t: &Token| -> bool { if let &Token::Junk = t { false } else { true } })
        .map(|t| t.string());
    let mut res = MarkovInfo::new();
    let mut last = Token::Begin.string();
    for token in str_tokens {
        res.add_token(last, &token);
        last = token.clone();
    }
    Ok(res)
}

pub fn make_beatnik(wottasquare: &str, markov_fname: &str) -> Result<String, io::Error> {
    let buffer = File::open(markov_fname).unwrap();
    let markov: MarkovInfo = serde_yaml::from_reader(&buffer).map_err(common::io_str_error)?;
    let words = runner::get_wottas(wottasquare)?;
    let mut last = Token::Begin.string();
    let mut out = String::new();
    for word in words {
        let token = markov.get_token(&last, word.score());
        out.push_str(&format!("{} ", token));
        last = token;
    }
    return Ok(out);
}

#[cfg(test)]
mod tests {
    use common;
    use quickcheck::TestResult;
    use super::{empty_filter, get_token};

    #[test]
    fn copes_with_utf_8() {
        common::get_words_core("why—I".as_bytes(), get_token, empty_filter).unwrap();
    }

    quickcheck! {
      fn token_test(xs: String) -> TestResult {
          return match common::get_words_core(xs.as_bytes(), get_token, empty_filter) {
              Ok(_) => TestResult::passed(),
              Err(err) => {
                  println!("Error: '{}'", xs);
                  TestResult::error(format!("{:?}", err))
              }
          }
      }
  }
}
<|endoftext|>


<commit_before>pub mod responses {

    use std::fmt;

    use serde;
    use serde_json as json;

    #[derive(Deserialize, Debug)]
    pub struct AuthResponse {
        pub session: SessionResponse,
    }

    #[derive(Deserialize, Debug, Clone)]
    pub struct SessionResponse {
        pub key: String,
        pub subscriber: i64,
        pub name: String,
    }

    #[derive(Deserialize)]
    pub struct NowPlayingResponseWrapper {
        pub nowplaying: NowPlayingResponse,
    }

    #[derive(Deserialize, Debug)]
    pub struct NowPlayingResponse {
        pub artist: CorrectableString,
        pub album: CorrectableString,
        #[serde(rename="albumArtist")]
        pub album_artist: CorrectableString,
        pub track: CorrectableString,
    }

    #[derive(Deserialize)]
    pub struct ScrobbleResponseWrapper {
        pub scrobbles: Scrobbles,
    }

    #[derive(Deserialize)]
    pub struct Scrobbles {
        pub scrobble: ScrobbleResponse,
    }

    #[derive(Deserialize, Debug)]
    pub struct ScrobbleResponse {
        pub artist: CorrectableString,
        pub album: CorrectableString,
        #[serde(rename="albumArtist")]
        pub album_artist: CorrectableString,
        pub track: CorrectableString,
        pub timestamp: String,
    }

    #[derive(Deserialize, Debug)]
    pub struct CorrectableString {
        #[serde(deserialize_with="CorrectableString::deserialize_corrected_field")]
        pub corrected: bool,
        #[serde(rename="#text", default)]
        pub text: String,
    }

    impl CorrectableString {
        fn deserialize_corrected_field<'de, D>(de: D) -> Result<bool, D::Error>
            where D: serde::Deserializer<'de>
        {
            let deser_result: json::Value = try!(serde::Deserialize::deserialize(de));
            match deser_result {
                json::Value::String(ref s) if &*s == "1" => Ok(true),
                json::Value::String(ref s) if &*s == "0" => Ok(false),
                _ => Err(serde::de::Error::custom("Unexpected value")),
            }
        }
    }

    impl fmt::Display for CorrectableString {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(f, "{}", self.text)
        }
    }

}

pub mod metadata {

    use std::collections::HashMap;

    pub struct Scrobble {
        artist: String,
        track: String,
        album: String
    }

    impl Scrobble {

        pub fn new(artist: String, track: String, album: String) -> Scrobble {
            Scrobble{ artist: artist, track: track, album: album }
        }

        pub fn as_map(&self) -> HashMap<&str, String> {
            let mut params = HashMap::new();
            params.insert("track", self.track.clone());
            params.insert("artist", self.artist.clone());
            params.insert("album", self.album.clone());

            params
        }

    }

}<commit_msg>First pass of documentation on Scrobble struct<commit_after>pub mod responses {

    use std::fmt;

    use serde;
    use serde_json as json;

    #[derive(Deserialize, Debug)]
    pub struct AuthResponse {
        pub session: SessionResponse,
    }

    #[derive(Deserialize, Debug, Clone)]
    pub struct SessionResponse {
        pub key: String,
        pub subscriber: i64,
        pub name: String,
    }

    #[derive(Deserialize)]
    pub struct NowPlayingResponseWrapper {
        pub nowplaying: NowPlayingResponse,
    }

    #[derive(Deserialize, Debug)]
    pub struct NowPlayingResponse {
        pub artist: CorrectableString,
        pub album: CorrectableString,
        #[serde(rename="albumArtist")]
        pub album_artist: CorrectableString,
        pub track: CorrectableString,
    }

    #[derive(Deserialize)]
    pub struct ScrobbleResponseWrapper {
        pub scrobbles: Scrobbles,
    }

    #[derive(Deserialize)]
    pub struct Scrobbles {
        pub scrobble: ScrobbleResponse,
    }

    #[derive(Deserialize, Debug)]
    pub struct ScrobbleResponse {
        pub artist: CorrectableString,
        pub album: CorrectableString,
        #[serde(rename="albumArtist")]
        pub album_artist: CorrectableString,
        pub track: CorrectableString,
        pub timestamp: String,
    }

    #[derive(Deserialize, Debug)]
    pub struct CorrectableString {
        #[serde(deserialize_with="CorrectableString::deserialize_corrected_field")]
        pub corrected: bool,
        #[serde(rename="#text", default)]
        pub text: String,
    }

    impl CorrectableString {
        fn deserialize_corrected_field<'de, D>(de: D) -> Result<bool, D::Error>
            where D: serde::Deserializer<'de>
        {
            let deser_result: json::Value = try!(serde::Deserialize::deserialize(de));
            match deser_result {
                json::Value::String(ref s) if &*s == "1" => Ok(true),
                json::Value::String(ref s) if &*s == "0" => Ok(false),
                _ => Err(serde::de::Error::custom("Unexpected value")),
            }
        }
    }

    impl fmt::Display for CorrectableString {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(f, "{}", self.text)
        }
    }

}

pub mod metadata {

    use std::collections::HashMap;

    /// Repesents a single track play (aka a "scrobble")
    pub struct Scrobble {
        artist: String,
        track: String,
        album: String
    }

    impl Scrobble {

        /// Constructs a new Scrobble instance, representing a music track
        /// played in the past.
        pub fn new(artist: String, track: String, album: String) -> Scrobble {
            Scrobble{ artist: artist, track: track, album: album }
        }

        /// Converts the Scrobble metadata (track name, artist & album name)
        /// into a HashMap. Map keys are `"track"`, `"artist"` and `"album"`,
        /// respectively.
        pub fn as_map(&self) -> HashMap<&str, String> {
            let mut params = HashMap::new();
            params.insert("track", self.track.clone());
            params.insert("artist", self.artist.clone());
            params.insert("album", self.album.clone());

            params
        }

    }

}<|endoftext|>


<commit_before>use std::ops::{Add, Sub, Div};
use std::collections::HashMap;
use std::iter::FromIterator;
use std::iter::repeat;
use std::fmt::{Debug, Formatter, Error};
use std::cmp::Ordering::Equal;
use std::ops::Drop;
use std::cmp;

#[derive(Clone,)]
pub struct WordVec {
    pub word: String,
    pub count: u64,
    vec: Vec<f32>,
}

impl PartialEq for WordVec {
    fn eq(&self, other: &WordVec) -> bool {
        self.vec == other.vec
    }
}

impl Debug for WordVec {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        try!(self.word.fmt(fmt));
        try!(fmt.write_str(": "));
        try!(self.count.fmt(fmt));
        try!(fmt.write_str(", "));
        self.vec.iter().take(5).collect::<Vec<_>>().fmt(fmt)
    }
}

impl WordVec {

    fn dot_prod(&self, other: &WordVec) -> f32 {
        self.vec.iter()
                .zip(other.vec.iter())
                .map(|(x, y)| x * y)
                .fold(0.0, |x, y| x + y)
    }

    fn magnitude(&self) -> f32 {
        self.vec.iter()
                .map(|x| x * x)
                .fold(0.0, |x, y| x + y)
                .sqrt()
    }

    pub fn distance(&self, other: &WordVec) -> f32 {
        self.vec.iter().zip(other.vec.iter())
                .map(|(x, y)| (x * x - y * y).abs().sqrt())
                .fold(0.0, |x, y| x + y)
    }

    fn new(word: String, num_words: usize) -> WordVec {
        WordVec {
            word: word,
            count: 0,
            vec: repeat(0.0).take(num_words).collect(),
        }
    }

    fn normalize(&mut self) {
        let count = self.count as f32;
        for f in &mut self.vec {
            *f /= count;
        }
    }

    fn inc(&mut self, i: usize) {
        self.vec[i] += 1.0;
    }

}

impl Div<i32> for WordVec {
    type Output = WordVec;

    fn div(mut self, i: i32) -> WordVec {
        for f in &mut self.vec {
            *f /= i as f32;
        }

        self
    }
}

impl<'a> Add<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn add(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} + {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..newvec.vec.len() {
            newvec.vec[i] += other.vec[i];
        }

        newvec
    }
}

impl<'a> Sub<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn sub(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} - {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..(newvec.vec.len()) {
            newvec.vec[i] -= other.vec[i];
        }

        newvec
    }
}

pub struct LanguageModel {
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct LanguageModelBuilder {
    window_radius: usize,
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct WordAcceptor<'a> {
    sentence: Vec<Option<usize>>,
    builder: &'a mut LanguageModelBuilder,
}

impl LanguageModelBuilder {
    pub fn new(window_radius: usize, words: Vec<String>) -> LanguageModelBuilder {

        let word_vecs = words.iter()
                             .map(|s| WordVec::new(s.clone(), words.len()))
                             .collect();

        let words = HashMap::from_iter(
                        words.into_iter()
                             .enumerate()
                             .map(|(a, b)| (b, a)));

        LanguageModelBuilder {
            window_radius: window_radius,
            words: words,
            word_vecs: word_vecs,
        }
    }

    pub fn build(mut self) -> LanguageModel {
        for vec in self.word_vecs.iter_mut() {
            vec.normalize();
        }

        LanguageModel {
            words: self.words,
            word_vecs: self.word_vecs,
        }
    }

    pub fn new_sentence<'a>(&'a mut self) -> WordAcceptor<'a> {
        WordAcceptor {
            sentence: Vec::new(),
            builder: self,
        }
    }

    fn get_vec(&mut self, word: &str) -> &mut WordVec {
        let i = self.words[word];
        &mut self.word_vecs[i]
    }

    fn add_word(&mut self, from: usize, to: usize) {
        self.word_vecs[from].inc(to);
    }

    fn word_seen(&mut self, word: &str) {
        self.get_vec(word).count += 1;
    }
}

impl<'a> WordAcceptor<'a> {
    pub fn add_word(&mut self, word: &str) {
        let idx_opt = self.builder.words.get(word).map(|w| *w);

        if let Some(next_idx) = idx_opt {
            self.builder.word_vecs[next_idx].count += 1;
            self.sentence.push(Some(next_idx));
        }
        else {
            self.sentence.push(None);
        }
    }
}

impl<'a> Drop for WordAcceptor<'a> {
    fn drop(&mut self) {
        let WordAcceptor {
            ref sentence,
            ref mut builder
        } = *self;

        println!("{:?}", sentence);

        for (i, &word_idx) in self.sentence.iter().enumerate() {
            let start = cmp::max(0, i as isize- builder.window_radius as isize) as usize;
            let end = cmp::min(sentence.len(), i + builder.window_radius + 1);
            for &j in &sentence[start..end] {
                match (word_idx, j) {
                    (Some(w), Some(j)) if w != j =>
                        builder.word_vecs[w].inc(j),
                    _ => {},
                }
            }
        }

        drop(sentence); // not actually sure if this is necessary
    }
}

impl LanguageModel {
    pub fn get(&self, word: &str) -> Option<&WordVec> {
        self.words.get(word).map(|i| &self.word_vecs[*i])
    }

    pub fn nearest_words(&self, word: &WordVec) -> Vec<&WordVec> {
        let mut vec_refs = self.word_vecs.iter().filter_map(|w| {
            let dist = w.distance(word);
            if w.word != word.word {
                Some((dist, w))
            } else {
                None
            }
        }).collect::<Vec<_>>();
        vec_refs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or_else(|| {
            println!("bad vector: {:?}, {:?}", a, b);
            Equal
        }));
        vec_refs.into_iter().map(|x| x.1).collect()
    }
}

#[cfg(test)]
mod test {
    use super::{WordVec, LanguageModel, LanguageModelBuilder};

    #[test]
    fn test_accept_sentences() {
        let words = "foo bar baz blort".words().map(|w| w.to_string()).collect::<Vec<_>>();
        let mut builder = LanguageModelBuilder::new(1, words);

        let input = "x foo bar baz x x x x x x x blort".words();

        {
            let mut acc = builder.new_sentence();
            for word in input {
                acc.add_word(word);
            }
        }

        let model = builder.build();
        let foo = model.get("foo").unwrap();
        let baz = model.get("baz").unwrap();
        let bar = model.get("bar").unwrap();
        let blort = model.get("blort").unwrap();

        assert!(foo.distance(baz) < foo.distance(blort));
        assert!(foo.distance(bar) == bar.distance(baz));

    }
}<commit_msg>Remove some cruft and fix calculation of distance.<commit_after>use std::ops::{Add, Sub, Div};
use std::collections::HashMap;
use std::iter::FromIterator;
use std::iter::repeat;
use std::fmt::{Debug, Formatter, Error};
use std::cmp::Ordering::Equal;
use std::ops::Drop;
use std::cmp;

#[derive(Clone,)]
pub struct WordVec {
    pub word: String,
    pub count: u64,
    vec: Vec<f32>,
}

impl PartialEq for WordVec {
    fn eq(&self, other: &WordVec) -> bool {
        self.vec == other.vec
    }
}

impl Debug for WordVec {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        try!(self.word.fmt(fmt));
        try!(fmt.write_str(": "));
        try!(self.count.fmt(fmt));
        try!(fmt.write_str(", "));
        self.vec.iter().take(5).collect::<Vec<_>>().fmt(fmt)
    }
}

impl WordVec {

    fn dot_prod(&self, other: &WordVec) -> f32 {
        self.vec.iter()
                .zip(other.vec.iter())
                .map(|(x, y)| x * y)
                .fold(0.0, |x, y| x + y)
    }

    fn magnitude(&self) -> f32 {
        self.vec.iter()
                .map(|x| x * x)
                .fold(0.0, |x, y| x + y)
                .sqrt()
    }

    pub fn distance(&self, other: &WordVec) -> f32 {
        self.vec.iter().zip(other.vec.iter())
                .map(|(x, y)| (x - y)*(x - y))
                .fold(0.0, |x, y| x + y)
                .abs().sqrt()
    }

    fn new(word: String, num_words: usize) -> WordVec {
        WordVec {
            word: word,
            count: 0,
            vec: repeat(0.0).take(num_words).collect(),
        }
    }

    fn normalize(&mut self) {
        let count = self.count as f32;
        for f in &mut self.vec {
            *f /= count;
        }
    }

    fn inc(&mut self, i: usize) {
        self.vec[i] += 1.0;
    }

}

impl Div<i32> for WordVec {
    type Output = WordVec;

    fn div(mut self, i: i32) -> WordVec {
        for f in &mut self.vec {
            *f /= i as f32;
        }

        self
    }
}

impl<'a> Add<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn add(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} + {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..newvec.vec.len() {
            newvec.vec[i] += other.vec[i];
        }

        newvec
    }
}

impl<'a> Sub<&'a WordVec> for WordVec {
    type Output = WordVec;

    fn sub(self, other: &WordVec) -> WordVec {
        debug_assert!(self.vec.len() == other.vec.len());
        let mut newvec = WordVec {
            word: format!("{} - {}", &self.word, &other.word),
            count: 0,
            vec: self.vec,
        };

        for i in 0..(newvec.vec.len()) {
            newvec.vec[i] -= other.vec[i];
        }

        newvec
    }
}

pub struct LanguageModel {
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct LanguageModelBuilder {
    window_radius: usize,
    words: HashMap<String, usize>,
    word_vecs: Vec<WordVec>,
}

pub struct WordAcceptor<'a> {
    sentence: Vec<Option<usize>>,
    builder: &'a mut LanguageModelBuilder,
}

impl LanguageModelBuilder {
    pub fn new(window_radius: usize, words: Vec<String>) -> LanguageModelBuilder {

        let word_vecs = words.iter()
                             .map(|s| WordVec::new(s.clone(), words.len()))
                             .collect();

        let words = HashMap::from_iter(
                        words.into_iter()
                             .enumerate()
                             .map(|(a, b)| (b, a)));

        LanguageModelBuilder {
            window_radius: window_radius,
            words: words,
            word_vecs: word_vecs,
        }
    }

    pub fn build(mut self) -> LanguageModel {
        for vec in self.word_vecs.iter_mut() {
            vec.normalize();
        }

        LanguageModel {
            words: self.words,
            word_vecs: self.word_vecs,
        }
    }

    pub fn new_sentence<'a>(&'a mut self) -> WordAcceptor<'a> {
        WordAcceptor {
            sentence: Vec::new(),
            builder: self,
        }
    }

    fn word_seen(&mut self, word: usize) {
        self.word_vecs[word].count += 1;
    }
}

impl<'a> WordAcceptor<'a> {
    pub fn add_word(&mut self, word: &str) {
        let idx_opt = self.builder.words.get(word).map(|w| *w);

        if let Some(next_idx) = idx_opt {
            self.builder.word_seen(next_idx);
            self.sentence.push(Some(next_idx));
        }
        else {
            self.sentence.push(None);
        }
    }
}

impl<'a> Drop for WordAcceptor<'a> {
    fn drop(&mut self) {
        let WordAcceptor {
            ref sentence,
            ref mut builder
        } = *self;

        for (i, &word_idx) in self.sentence.iter().enumerate() {
            let start = cmp::max(0, i as isize- builder.window_radius as isize) as usize;
            let end = cmp::min(sentence.len(), i + builder.window_radius + 1);
            for &j in &sentence[start..end] {
                match (word_idx, j) {
                    (Some(w), Some(j)) if w != j =>
                        builder.word_vecs[w].inc(j),
                    _ => {},
                }
            }
        }

        drop(sentence); // not actually sure if this is necessary
    }
}

impl LanguageModel {
    pub fn get(&self, word: &str) -> Option<&WordVec> {
        self.words.get(word).map(|i| &self.word_vecs[*i])
    }

    pub fn nearest_words(&self, word: &WordVec) -> Vec<&WordVec> {
        let mut vec_refs = self.word_vecs.iter().filter_map(|w| {
            let dist = w.distance(word);
            if w.word != word.word {
                Some((dist, w))
            } else {
                None
            }
        }).collect::<Vec<_>>();
        vec_refs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or_else(|| {
            println!("bad vector: {:?}, {:?}", a, b);
            Equal
        }));
        vec_refs.into_iter().map(|x| x.1).collect()
    }
}

#[cfg(test)]
mod test {
    use super::{WordVec, LanguageModel, LanguageModelBuilder};

    #[test]
    fn test_accept_sentences() {
        let words = "foo bar baz blort".words().map(|w| w.to_string()).collect::<Vec<_>>();
        let mut builder = LanguageModelBuilder::new(1, words);

        let input = "x foo bar baz x x x x x x x blort".words();

        {
            let mut acc = builder.new_sentence();
            for word in input {
                acc.add_word(word);
            }
        }

        let model = builder.build();
        let foo = model.get("foo").unwrap();
        let baz = model.get("baz").unwrap();
        let bar = model.get("bar").unwrap();
        let blort = model.get("blort").unwrap();

        assert!(foo.distance(baz) < foo.distance(blort));
        assert!(foo.distance(bar) == bar.distance(baz));

    }
}<|endoftext|>