<commit_before>use irc::parse::IrcMsg;
use rand::{thread_rng, Rng};

use command_mapper::{
    RustBotPlugin,
    CommandMapperDispatch,
    IrcBotConfigurator,
    Format,
    Token,
};

const CMD_8BALL_QUERY: Token = Token(0);
const CMD_8BALL: Token = Token(1);

const ANSWERS: &'static [&'static str] = &[
    "It is certain",
    "It is decidedly so",
    "Without a doubt",
    "Yes, definitely",
    "You may rely on it",
    "As I see it, yes",
    "Most likely",
    "Outlook good",
    "Yes",
    "Signs point to yes",
    "Reply hazy try again",
    "Ask again later",
    "Better not tell you now",
    "Cannot predict now",
    "Concentrate and ask again",
    "Don't count on it",
    "My reply is no",
    "My sources say no",
    "Outlook not so good",
    "Very doubtful",
];

pub struct EightBallPlugin;

impl EightBallPlugin {
    pub fn new() -> EightBallPlugin {
        EightBallPlugin
    }

    pub fn get_plugin_name() -> &'static str {
        "8ball"
    }
}

enum EightBallCommandType {
    EightBall,
}

fn parse_command<'a>(m: &CommandMapperDispatch) -> Option<EightBallCommandType> {
    match m.command().token {
        CMD_8BALL_QUERY => Some(EightBallCommandType::EightBall),
        CMD_8BALL => Some(EightBallCommandType::EightBall),
        _ => None
    }
}

impl RustBotPlugin for EightBallPlugin {
    fn configure(&mut self, conf: &mut IrcBotConfigurator) {
        conf.map_format(CMD_8BALL_QUERY, Format::from_str("8ball {*query}").ok().unwrap());
        conf.map_format(CMD_8BALL, Format::from_str("8ball").ok().unwrap());
    }

    fn dispatch_cmd(&mut self, m: &CommandMapperDispatch, msg: &IrcMsg) {
        match parse_command(m) {
            Some(EightBallCommandType::EightBall) => {
                let prefix = msg.get_prefix();
                let nick = match prefix.nick() {
                    Some(nick) => nick,
                    None => return,
                };
                let answer = match thread_rng().choose(ANSWERS) {
                    Some(answer) => answer,
                    None => return,
                };
                let msg = format!("{}: {}", nick, answer);
                m.reply(&msg);
            },
            None => return
        }
    }
}
<commit_msg>plugin(8ball): Unicodify<commit_after>use irc::parse::IrcMsg;
use rand::{thread_rng, Rng};

use command_mapper::{
    RustBotPlugin,
    CommandMapperDispatch,
    IrcBotConfigurator,
    Format,
    Token,
};

const CMD_8BALL: Token = Token(1);

const ANSWERS: &'static [&'static str] = &[
    "It is certain",
    "It is decidedly so",
    "Without a doubt",
    "Yes, definitely",
    "You may rely on it",
    "As I see it, yes",
    "Most likely",
    "Outlook good",
    "Yes",
    "Signs point to yes",
    "Reply hazy try again",
    "Ask again later",
    "Better not tell you now",
    "Cannot predict now",
    "Concentrate and ask again",
    "Don't count on it",
    "My reply is no",
    "My sources say no",
    "Outlook not so good",
    "Very doubtful",
];

pub struct EightBallPlugin;

impl EightBallPlugin {
    pub fn new() -> EightBallPlugin {
        EightBallPlugin
    }

    pub fn get_plugin_name() -> &'static str {
        "8ball"
    }
}

enum EightBallCommandType {
    EightBall,
}

fn parse_command<'a>(m: &CommandMapperDispatch) -> Option<EightBallCommandType> {
    match m.command().token {
        CMD_8BALL => Some(EightBallCommandType::EightBall),
        _ => None
    }
}

impl RustBotPlugin for EightBallPlugin {
    fn configure(&mut self, conf: &mut IrcBotConfigurator) {
        conf.map_format(CMD_8BALL, Format::from_str("8ball {*query}").ok().unwrap());
        conf.map_format(CMD_8BALL, Format::from_str("8ball").ok().unwrap());
        conf.map_format(CMD_8BALL, Format::from_str("\u{1f3b1} {*query}").ok().unwrap());
        conf.map_format(CMD_8BALL, Format::from_str("\u{1f3b1}").ok().unwrap());
    }

    fn dispatch_cmd(&mut self, m: &CommandMapperDispatch, msg: &IrcMsg) {
        match parse_command(m) {
            Some(EightBallCommandType::EightBall) => {
                let prefix = msg.get_prefix();
                let nick = match prefix.nick() {
                    Some(nick) => nick,
                    None => return,
                };
                let answer = match thread_rng().choose(ANSWERS) {
                    Some(answer) => answer,
                    None => return,
                };
                let msg = format!("{}: {}", nick, answer);
                m.reply(&msg);
            },
            None => return
        }
    }
}
<|endoftext|>


<commit_before>//! Defines the Whitted integrator which implements Whitted recursive ray tracing

use std::f32;
use rand::StdRng;

use scene::Scene;
use linalg::{self, Ray};
use geometry::{Intersection, Emitter};
use film::Colorf;
use integrator::Integrator;
use bxdf::BxDFType;
use light::Light;
use sampler::Sampler;

/// The Whitted integrator implementing the Whitted recursive ray tracing algorithm
/// See [Whitted, An improved illumination model for shaded display](http://dl.acm.org/citation.cfm?id=358882)
#[derive(Clone, Copy, Debug)]
pub struct Whitted {
    /// The maximum recursion depth for rays
    max_depth: u32,
}

impl Whitted {
    /// Create a new Whitted integrator with the desired maximum recursion depth for rays
    pub fn new(max_depth: u32) -> Whitted { Whitted { max_depth: max_depth } }
}

impl Integrator for Whitted {
    fn illumination(&self, scene: &Scene, light_list: &Vec<&Emitter>, ray: &Ray,
                    hit: &Intersection, sampler: &mut Sampler, rng: &mut StdRng) -> Colorf {
        let bsdf = hit.material.bsdf(hit);
        let w_o = -ray.d;
        let mut sample_2d = [(0.0, 0.0)];
        sampler.get_samples_2d(&mut sample_2d[..], rng);
        let mut illum = Colorf::broadcast(0.0);
        for light in light_list {
            let (li, w_i, pdf, occlusion) = light.sample_incident(&hit.dg.p, &sample_2d[0]);
            let f = bsdf.eval(&w_o, &w_i, BxDFType::all());
            if !li.is_black() && !f.is_black() && !occlusion.occluded(scene) {
                illum = illum + f * li * f32::abs(linalg::dot(&w_i, &bsdf.n)) / pdf;
            }
        }
        if ray.depth < self.max_depth {
            illum = illum + self.specular_reflection(scene, light_list, ray, &bsdf, sampler, rng);
            illum = illum + self.specular_transmission(scene, light_list, ray, &bsdf, sampler, rng);
        }
        illum
    }
}

<commit_msg>Have whitted rendered include emission of directly hit lights<commit_after>//! Defines the Whitted integrator which implements Whitted recursive ray tracing

use std::f32;
use rand::StdRng;

use scene::Scene;
use linalg::{self, Ray};
use geometry::{Intersection, Emitter, Instance};
use film::Colorf;
use integrator::Integrator;
use bxdf::BxDFType;
use light::Light;
use sampler::Sampler;

/// The Whitted integrator implementing the Whitted recursive ray tracing algorithm
/// See [Whitted, An improved illumination model for shaded display](http://dl.acm.org/citation.cfm?id=358882)
#[derive(Clone, Copy, Debug)]
pub struct Whitted {
    /// The maximum recursion depth for rays
    max_depth: u32,
}

impl Whitted {
    /// Create a new Whitted integrator with the desired maximum recursion depth for rays
    pub fn new(max_depth: u32) -> Whitted { Whitted { max_depth: max_depth } }
}

impl Integrator for Whitted {
    fn illumination(&self, scene: &Scene, light_list: &Vec<&Emitter>, ray: &Ray,
                    hit: &Intersection, sampler: &mut Sampler, rng: &mut StdRng) -> Colorf {
        let bsdf = hit.material.bsdf(hit);
        let w_o = -ray.d;
        let mut sample_2d = [(0.0, 0.0)];
        sampler.get_samples_2d(&mut sample_2d[..], rng);
        let mut illum = Colorf::broadcast(0.0);
        if let &Instance::Emitter(ref e) = hit.instance {
            let w = -ray.d;
            illum = illum + e.radiance(&w, &hit.dg.p, &hit.dg.ng);
        }

        for light in light_list {
            let (li, w_i, pdf, occlusion) = light.sample_incident(&hit.dg.p, &sample_2d[0]);
            let f = bsdf.eval(&w_o, &w_i, BxDFType::all());
            if !li.is_black() && !f.is_black() && !occlusion.occluded(scene) {
                illum = illum + f * li * f32::abs(linalg::dot(&w_i, &bsdf.n)) / pdf;
            }
        }
        if ray.depth < self.max_depth {
            illum = illum + self.specular_reflection(scene, light_list, ray, &bsdf, sampler, rng);
            illum = illum + self.specular_transmission(scene, light_list, ray, &bsdf, sampler, rng);
        }
        illum
    }
}

<|endoftext|>


<commit_before>#![crate_name = "ansi_term"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(phase, macro_rules)]

//! This is a library for controlling colours and formatting, such as
//! red bold text or blue underlined text, on ANSI terminals.
//!
//! ```rust
//! extern crate ansi_term;
//! use ansi_term::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, Fixed};
//! use ansi_term::Style::Plain;
//! ```
//!
//! Simple Colours
//! --------------
//!
//! You can format strings by calling the `paint` method on a Colour
//! or a Style object, passing in the string you want to format. For
//! example, to get some red text, call the `paint` method on `Red`:
//!
//! ```rust
//! println!("This is in red: {}!", Red.paint("a red string"));
//! ```
//!
//! The `paint` method returns an `ANSIString` object, which will get
//! automatically converted to the correct sequence of escape codes when
//! used in a `println!` or `format!` macro, or anything else that
//! supports using the `Show` trait. This means that if you just want a
//! string of the escape codes without anything else, you can still use
//! the `to_string` method:
//!
//! ```rust
//! let red_string: String = Red.paint("another red string").to_string();
//! ```
//!
//! Bold, Underline, and Background
//! -------------------------------
//!
//! To do anything more complex than just foreground colours, you need
//! to use Style objects. Calling the `bold` or `underline` method on
//! a Colour returns a Style that has the appropriate property set on
//! it:
//!
//! ```rust
//! println!("Demonstrating {} and {}!",
//!          Blue.bold().paint("blue bold"),
//!          Yellow.underline().paint("yellow underline"));
//! ```
//!
//! These methods chain, so you can call them on existing Style
//! objects to set more than one particular properly, like so:
//!
//! ```rust
//! Blue.underline().bold().paint("Blue underline bold!")
//! ```
//!
//! You can set the background colour of a Style by using the `on`
//! method:
//!
//! ```rust
//! Blue.on(Yellow).paint("Blue on yellow!")
//! ```
//!
//! Finally, you can turn a Colour into a Style with the `normal`
//! method, though it'll produce the exact same string if you just use
//! the Colour. It's only useful if you're writing a method that can
//! return either normal or bold (or underline) styles, and need to
//! return a Style object from it.
//!
//! ```rust
//! Red.normal().paint("yet another red string")
//! ```
//!
//! Extended Colours
//! ----------------
//!
//! You can access the extended range of 256 colours by using the
//! Fixed constructor, which takes an argument of the colour number to
//! use. This can be used wherever you would use a Colour:
//!
//! ```rust
//! Fixed(134).paint("A sort of light purple.")
//! ```
//!
//! This even works for background colours:
//!
//! ```rust
//! Fixed(221).on(Fixed(124)).paint("Mustard in the ketchup.")
//! ```
//!
//! No Formatting
//! -------------
//!
//! Finally, for the sake of completeness, the Plain style provides
//! neither colours nor formatting.
//!
//! ```rust
//! Plain.paint("No colours here.")
//! ```

#[phase(plugin)]
extern crate regex_macros;
extern crate regex;

use Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
use Style::{Plain, Foreground, Styled};
use std::fmt;

/// An ANSI String is a string coupled with the Style to display it
/// in a terminal.
///
/// Although not technically a string itself, it can be turned into
/// one with the `to_string` method.
pub struct ANSIString<'a> {
    string: &'a str,
    style: Style,
}

impl<'a> Copy for ANSIString<'a> { }

impl<'a> ANSIString<'a> {
    /// Creates a new ANSI String with the given contents and style.
    pub fn new(contents: &'a str, style: Style) -> ANSIString<'a> {
        ANSIString { string: contents, style: style }
    }
}

impl<'a> fmt::Show for ANSIString<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.style {
            Plain => write!(f, "{}", self.string),
            Foreground(colour) => {
                try!(f.write_str("\x1B["));
                try!(colour.write_foreground_code(f));
                write!(f, "m{}\x1B[0m", self.string)
            },
            Styled { foreground, background, bold, underline } => {
                try!(f.write_str("\x1B["));

                if bold {
                    try!(f.write_str("1;"));
                }

                if underline {
                    try!(f.write_str("4;"));
                }

                match background {
                    Some(c) => {
                        try!(c.write_background_code(f));
                        try!(f.write_str(";"));
                    },
                    None => {},
                }

                try!(foreground.write_foreground_code(f));

                write!(f, "m{}\x1B[0m", self.string)
            }
        }
    }
}

/// A colour is one specific type of ANSI escape code, and can refer
/// to either the foreground or background colour.
///
/// These use the standard numeric sequences.
/// See http://invisible-island.net/xterm/ctlseqs/ctlseqs.html
pub enum Colour {
    Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed(u8),
}

impl Copy for Colour { }

// I'm not beyond calling Colour Colour, rather than Color, but I did
// purposefully name this crate 'ansi-term' so people wouldn't get
// confused when they tried to install it.
//
// Only *after* they'd installed it.

impl Colour {
    fn write_foreground_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("30"),
            Red    => f.write_str("31"),
            Green  => f.write_str("32"),
            Yellow => f.write_str("33"),
            Blue   => f.write_str("34"),
            Purple => f.write_str("35"),
            Cyan   => f.write_str("36"),
            White  => f.write_str("37"),
            Fixed(num) => write!(f, "38;5;{}", num),
        }
    }

    fn write_background_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("40"),
            Red    => f.write_str("41"),
            Green  => f.write_str("42"),
            Yellow => f.write_str("43"),
            Blue   => f.write_str("44"),
            Purple => f.write_str("45"),
            Cyan   => f.write_str("46"),
            White  => f.write_str("47"),
            Fixed(num) => write!(f, "48;5;{}", num),
        }
    }

    /// Return a Style with the foreground colour set to this colour.
    pub fn normal(self) -> Style {
        Styled { foreground: self, background: None, bold: false, underline: false }
    }

    /// Paints the given text with this colour, returning an ANSI string.
    /// This is a short-cut so you don't have to use Blue.normal() just
    /// to get blue text.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, Foreground(self))
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        Styled { foreground: self, background: None, bold: false, underline: true }
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        Styled { foreground: self, background: None, bold: true, underline: false }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        Styled { foreground: self, background: Some(background), bold: false, underline: false }
    }
}

/// A style is a collection of properties that can format a string
/// using ANSI escape codes.
pub enum Style {

    /// The Plain style provides no formatting.
    Plain,

    /// The Foreground style provides just a foreground colour.
    Foreground(Colour),

    /// The Styled style is a catch-all for anything more complicated
    /// than that. It's technically possible for there to be other
    /// cases, such as "bold foreground", but probably isn't worth it.
    Styled { foreground: Colour, background: Option<Colour>, bold: bool, underline: bool, },
}

impl Copy for Style { }

impl Style {
    /// Paints the given text with this colour, returning an ANSI string.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, self)
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        match self {
            Plain => Styled { foreground: White, background: None, bold: true, underline: false },
            Foreground(c) => Styled { foreground: c, background: None, bold: true, underline: false },
            Styled { foreground, background, bold: _, underline } => Styled { foreground: foreground, background: background, bold: true, underline: underline },
        }
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        match self {
            Plain => Styled { foreground: White, background: None, bold: false, underline: true },
            Foreground(c) => Styled { foreground: c, background: None, bold: false, underline: true },
            Styled { foreground, background, bold, underline: _ } => Styled { foreground: foreground, background: background, bold: bold, underline: true },
        }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        match self {
            Plain => Styled { foreground: White,background: Some(background), bold: false, underline: false },
            Foreground(c) => Styled { foreground: c, background: Some(background), bold: false, underline: false },
            Styled { foreground, background: _, bold, underline } => Styled { foreground: foreground, background: Some(background), bold: bold, underline: underline },
        }
    }
}

/// Return a String with all ANSI escape codes removed.
///
/// For example:
///
/// ```rust
/// strip_formatting(Blue.paint("hello!").to_string()) == "hello!"
/// ```
pub fn strip_formatting(input: &str) -> String {
    // What's blue and smells like red paint? Blue paint.
    let re = regex!("\x1B\\[.+?m");
    re.replace_all(input, "").to_string()
}

#[cfg(test)]
mod tests {
    use super::strip_formatting;
    use super::Style::Plain;
    use super::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};

    macro_rules! test {
        ($name: ident: $style: expr $input: expr => $result: expr) => {
            #[test]
            fn $name() {
                assert_eq!($style.paint($input).to_string(), $result.to_string())
            }
        };
    }

    test!(plain:                 Plain                             "text/plain" => "text/plain");
    test!(red:                   Red                               "hi" => "\x1B[31mhi\x1B[0m");
    test!(black:                 Black.normal()                    "hi" => "\x1B[30mhi\x1B[0m");
    test!(yellow_bold:           Yellow.bold()                     "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(yellow_bold_2:         Yellow.normal().bold()            "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(blue_underline:        Blue.underline()                  "hi" => "\x1B[4;34mhi\x1B[0m");
    test!(green_bold_ul:         Green.bold().underline()          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(green_bold_ul_2:       Green.underline().bold()          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(purple_on_white:       Purple.on(White)                  "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(purple_on_white_2:     Purple.normal().on(White)         "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(cyan_bold_on_white:    Cyan.bold().on(White)             "hi" => "\x1B[1;47;36mhi\x1B[0m");
    test!(cyan_ul_on_white:      Cyan.underline().on(White)        "hi" => "\x1B[4;47;36mhi\x1B[0m");
    test!(cyan_bold_ul_on_white: Cyan.bold().underline().on(White) "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(cyan_ul_bold_on_white: Cyan.underline().bold().on(White) "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(fixed:                 Fixed(100)                        "hi" => "\x1B[38;5;100mhi\x1B[0m");
    test!(fixed_on_purple:       Fixed(100).on(Purple)             "hi" => "\x1B[45;38;5;100mhi\x1B[0m");
    test!(fixed_on_fixed:        Fixed(100).on(Fixed(200))         "hi" => "\x1B[48;5;200;38;5;100mhi\x1B[0m");

    #[test]
    fn test_strip_formatting() {
        let hi = strip_formatting(Blue.paint("hi").to_string().as_slice());
        assert!(hi == "hi".to_string());
    }

    #[test]
    fn test_strip_formatting_2() {
        let hi = strip_formatting(Blue.on(Fixed(230)).bold().paint("hi").to_string().as_slice());
        assert!(hi == "hi".to_string());
    }
}
<commit_msg>Upgrade to latest Rust<commit_after>#![crate_name = "ansi_term"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(plugin)]

//! This is a library for controlling colours and formatting, such as
//! red bold text or blue underlined text, on ANSI terminals.
//!
//! ```rust
//! extern crate ansi_term;
//! use ansi_term::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, Fixed};
//! use ansi_term::Style::Plain;
//! ```
//!
//! Simple Colours
//! --------------
//!
//! You can format strings by calling the `paint` method on a Colour
//! or a Style object, passing in the string you want to format. For
//! example, to get some red text, call the `paint` method on `Red`:
//!
//! ```rust
//! println!("This is in red: {}!", Red.paint("a red string"));
//! ```
//!
//! The `paint` method returns an `ANSIString` object, which will get
//! automatically converted to the correct sequence of escape codes when
//! used in a `println!` or `format!` macro, or anything else that
//! supports using the `Show` trait. This means that if you just want a
//! string of the escape codes without anything else, you can still use
//! the `to_string` method:
//!
//! ```rust
//! let red_string: String = Red.paint("another red string").to_string();
//! ```
//!
//! Bold, Underline, and Background
//! -------------------------------
//!
//! To do anything more complex than just foreground colours, you need
//! to use Style objects. Calling the `bold` or `underline` method on
//! a Colour returns a Style that has the appropriate property set on
//! it:
//!
//! ```rust
//! println!("Demonstrating {} and {}!",
//!          Blue.bold().paint("blue bold"),
//!          Yellow.underline().paint("yellow underline"));
//! ```
//!
//! These methods chain, so you can call them on existing Style
//! objects to set more than one particular properly, like so:
//!
//! ```rust
//! Blue.underline().bold().paint("Blue underline bold!")
//! ```
//!
//! You can set the background colour of a Style by using the `on`
//! method:
//!
//! ```rust
//! Blue.on(Yellow).paint("Blue on yellow!")
//! ```
//!
//! Finally, you can turn a Colour into a Style with the `normal`
//! method, though it'll produce the exact same string if you just use
//! the Colour. It's only useful if you're writing a method that can
//! return either normal or bold (or underline) styles, and need to
//! return a Style object from it.
//!
//! ```rust
//! Red.normal().paint("yet another red string")
//! ```
//!
//! Extended Colours
//! ----------------
//!
//! You can access the extended range of 256 colours by using the
//! Fixed constructor, which takes an argument of the colour number to
//! use. This can be used wherever you would use a Colour:
//!
//! ```rust
//! Fixed(134).paint("A sort of light purple.")
//! ```
//!
//! This even works for background colours:
//!
//! ```rust
//! Fixed(221).on(Fixed(124)).paint("Mustard in the ketchup.")
//! ```
//!
//! No Formatting
//! -------------
//!
//! Finally, for the sake of completeness, the Plain style provides
//! neither colours nor formatting.
//!
//! ```rust
//! Plain.paint("No colours here.")
//! ```

#[plugin]
extern crate regex_macros;
extern crate regex;

use Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
use Style::{Plain, Foreground, Styled};
use std::fmt;

/// An ANSI String is a string coupled with the Style to display it
/// in a terminal.
///
/// Although not technically a string itself, it can be turned into
/// one with the `to_string` method.
pub struct ANSIString<'a> {
    string: &'a str,
    style: Style,
}

impl<'a> Copy for ANSIString<'a> { }

impl<'a> ANSIString<'a> {
    /// Creates a new ANSI String with the given contents and style.
    pub fn new(contents: &'a str, style: Style) -> ANSIString<'a> {
        ANSIString { string: contents, style: style }
    }
}

impl<'a> fmt::Show for ANSIString<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.style {
            Plain => write!(f, "{}", self.string),
            Foreground(colour) => {
                try!(f.write_str("\x1B["));
                try!(colour.write_foreground_code(f));
                write!(f, "m{}\x1B[0m", self.string)
            },
            Styled { foreground, background, bold, underline } => {
                try!(f.write_str("\x1B["));

                if bold {
                    try!(f.write_str("1;"));
                }

                if underline {
                    try!(f.write_str("4;"));
                }

                match background {
                    Some(c) => {
                        try!(c.write_background_code(f));
                        try!(f.write_str(";"));
                    },
                    None => {},
                }

                try!(foreground.write_foreground_code(f));

                write!(f, "m{}\x1B[0m", self.string)
            }
        }
    }
}

/// A colour is one specific type of ANSI escape code, and can refer
/// to either the foreground or background colour.
///
/// These use the standard numeric sequences.
/// See http://invisible-island.net/xterm/ctlseqs/ctlseqs.html
pub enum Colour {
    Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed(u8),
}

impl Copy for Colour { }

// I'm not beyond calling Colour Colour, rather than Color, but I did
// purposefully name this crate 'ansi-term' so people wouldn't get
// confused when they tried to install it.
//
// Only *after* they'd installed it.

impl Colour {
    fn write_foreground_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("30"),
            Red    => f.write_str("31"),
            Green  => f.write_str("32"),
            Yellow => f.write_str("33"),
            Blue   => f.write_str("34"),
            Purple => f.write_str("35"),
            Cyan   => f.write_str("36"),
            White  => f.write_str("37"),
            Fixed(num) => write!(f, "38;5;{}", num),
        }
    }

    fn write_background_code(self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Black  => f.write_str("40"),
            Red    => f.write_str("41"),
            Green  => f.write_str("42"),
            Yellow => f.write_str("43"),
            Blue   => f.write_str("44"),
            Purple => f.write_str("45"),
            Cyan   => f.write_str("46"),
            White  => f.write_str("47"),
            Fixed(num) => write!(f, "48;5;{}", num),
        }
    }

    /// Return a Style with the foreground colour set to this colour.
    pub fn normal(self) -> Style {
        Styled { foreground: self, background: None, bold: false, underline: false }
    }

    /// Paints the given text with this colour, returning an ANSI string.
    /// This is a short-cut so you don't have to use Blue.normal() just
    /// to get blue text.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, Foreground(self))
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        Styled { foreground: self, background: None, bold: false, underline: true }
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        Styled { foreground: self, background: None, bold: true, underline: false }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        Styled { foreground: self, background: Some(background), bold: false, underline: false }
    }
}

/// A style is a collection of properties that can format a string
/// using ANSI escape codes.
pub enum Style {

    /// The Plain style provides no formatting.
    Plain,

    /// The Foreground style provides just a foreground colour.
    Foreground(Colour),

    /// The Styled style is a catch-all for anything more complicated
    /// than that. It's technically possible for there to be other
    /// cases, such as "bold foreground", but probably isn't worth it.
    Styled { foreground: Colour, background: Option<Colour>, bold: bool, underline: bool, },
}

impl Copy for Style { }

impl Style {
    /// Paints the given text with this colour, returning an ANSI string.
    pub fn paint(self, input: &str) -> ANSIString {
        ANSIString::new(input, self)
    }

    /// Returns a Style with the bold property set.
    pub fn bold(self) -> Style {
        match self {
            Plain => Styled { foreground: White, background: None, bold: true, underline: false },
            Foreground(c) => Styled { foreground: c, background: None, bold: true, underline: false },
            Styled { foreground, background, bold: _, underline } => Styled { foreground: foreground, background: background, bold: true, underline: underline },
        }
    }

    /// Returns a Style with the underline property set.
    pub fn underline(self) -> Style {
        match self {
            Plain => Styled { foreground: White, background: None, bold: false, underline: true },
            Foreground(c) => Styled { foreground: c, background: None, bold: false, underline: true },
            Styled { foreground, background, bold, underline: _ } => Styled { foreground: foreground, background: background, bold: bold, underline: true },
        }
    }

    /// Returns a Style with the background colour property set.
    pub fn on(self, background: Colour) -> Style {
        match self {
            Plain => Styled { foreground: White,background: Some(background), bold: false, underline: false },
            Foreground(c) => Styled { foreground: c, background: Some(background), bold: false, underline: false },
            Styled { foreground, background: _, bold, underline } => Styled { foreground: foreground, background: Some(background), bold: bold, underline: underline },
        }
    }
}

/// Return a String with all ANSI escape codes removed.
///
/// For example:
///
/// ```rust
/// strip_formatting(Blue.paint("hello!").to_string()) == "hello!"
/// ```
pub fn strip_formatting(input: &str) -> String {
    // What's blue and smells like red paint? Blue paint.
    let re = regex!("\x1B\\[.+?m");
    re.replace_all(input, "").to_string()
}

#[cfg(test)]
mod tests {
    use super::strip_formatting;
    use super::Style::Plain;
    use super::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};

    macro_rules! test {
        ($name: ident: $style: expr $input: expr => $result: expr) => {
            #[test]
            fn $name() {
                assert_eq!($style.paint($input).to_string(), $result.to_string())
            }
        };
    }

    test!(plain:                 Plain                             "text/plain" => "text/plain");
    test!(red:                   Red                               "hi" => "\x1B[31mhi\x1B[0m");
    test!(black:                 Black.normal()                    "hi" => "\x1B[30mhi\x1B[0m");
    test!(yellow_bold:           Yellow.bold()                     "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(yellow_bold_2:         Yellow.normal().bold()            "hi" => "\x1B[1;33mhi\x1B[0m");
    test!(blue_underline:        Blue.underline()                  "hi" => "\x1B[4;34mhi\x1B[0m");
    test!(green_bold_ul:         Green.bold().underline()          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(green_bold_ul_2:       Green.underline().bold()          "hi" => "\x1B[1;4;32mhi\x1B[0m");
    test!(purple_on_white:       Purple.on(White)                  "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(purple_on_white_2:     Purple.normal().on(White)         "hi" => "\x1B[47;35mhi\x1B[0m");
    test!(cyan_bold_on_white:    Cyan.bold().on(White)             "hi" => "\x1B[1;47;36mhi\x1B[0m");
    test!(cyan_ul_on_white:      Cyan.underline().on(White)        "hi" => "\x1B[4;47;36mhi\x1B[0m");
    test!(cyan_bold_ul_on_white: Cyan.bold().underline().on(White) "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(cyan_ul_bold_on_white: Cyan.underline().bold().on(White) "hi" => "\x1B[1;4;47;36mhi\x1B[0m");
    test!(fixed:                 Fixed(100)                        "hi" => "\x1B[38;5;100mhi\x1B[0m");
    test!(fixed_on_purple:       Fixed(100).on(Purple)             "hi" => "\x1B[45;38;5;100mhi\x1B[0m");
    test!(fixed_on_fixed:        Fixed(100).on(Fixed(200))         "hi" => "\x1B[48;5;200;38;5;100mhi\x1B[0m");

    #[test]
    fn test_strip_formatting() {
        let hi = strip_formatting(Blue.paint("hi").to_string().as_slice());
        assert!(hi == "hi".to_string());
    }

    #[test]
    fn test_strip_formatting_2() {
        let hi = strip_formatting(Blue.on(Fixed(230)).bold().paint("hi").to_string().as_slice());
        assert!(hi == "hi".to_string());
    }
}
<|endoftext|>


<commit_before>
use std::sync::Arc;
use std::collections::HashMap;
use std::borrow::Cow;
use parking_lot::RwLock;
use elements::{Internal, ValueType};
use interpreter::Error;
use interpreter::module::{ModuleInstanceInterface, ExecutionParams, ItemIndex,
	CallerContext, ExportEntryType, InternalFunctionReference, InternalFunction, FunctionSignature};
use interpreter::memory::MemoryInstance;
use interpreter::table::TableInstance;
use interpreter::value::RuntimeValue;
use interpreter::variable::{VariableInstance, VariableType};

/// Min index of native function.
pub const NATIVE_INDEX_FUNC_MIN: u32 = 10001;
/// Min index of native global.
pub const NATIVE_INDEX_GLOBAL_MIN: u32 = 20001;

/// User functions executor.
pub trait UserFunctionExecutor {
	/// Execute function with given name.
	fn execute(&mut self, name: &str, context: CallerContext) -> Result<Option<RuntimeValue>, Error>;
}

/// User function descriptor
#[derive(Debug, Clone)]
pub enum UserFunctionDescriptor {
	/// Static function definition
	Static(&'static str, &'static [ValueType], Option<ValueType>),
	/// Dynamic heap function definition
	Heap(String, Vec<ValueType>, Option<ValueType>),
}

impl UserFunctionDescriptor {
	/// New function with statically known params
	pub fn statik(name: &'static str, params: &'static [ValueType], result: Option<ValueType>) -> Self {
		UserFunctionDescriptor::Static(name, params, result)
	}

	/// New function with statically unknown params
	pub fn heap(name: String, params: Vec<ValueType>, result: Option<ValueType>) -> Self {
		UserFunctionDescriptor::Heap(name, params, result)
	}

	/// Name of the function
	pub fn name(&self) -> &str {
		match self {
			&UserFunctionDescriptor::Static(name, _, _) => name,
			&UserFunctionDescriptor::Heap(ref name, _, _) => name,
		}
	}

	/// Arguments of the function
	pub fn params(&self) -> &[ValueType] {
		match self {
			&UserFunctionDescriptor::Static(_, params, _) => params,
			&UserFunctionDescriptor::Heap(_, ref params, _) => params,
		}
	}

	/// Return type of the function
	pub fn return_type(&self) -> Option<ValueType> {
		match self {
			&UserFunctionDescriptor::Static(_, _, result) => result,
			&UserFunctionDescriptor::Heap(_, _, result) => result,
		}
	}
}

/// Set of user-defined module elements.
pub struct UserDefinedElements<E: UserFunctionExecutor> {
	/// User globals list.
	pub globals: HashMap<String, Arc<VariableInstance>>,
	/// User functions list.
	pub functions: Cow<'static, [UserFunctionDescriptor]>,
	/// Functions executor.
	pub executor: Option<E>,
}

/// Native module instance.
pub struct NativeModuleInstance<E: UserFunctionExecutor> {
	/// Underlying module reference.
	base: Arc<ModuleInstanceInterface>,
	/// User function executor.
	executor: RwLock<Option<E>>,
	/// By-name functions index.
	functions_by_name: HashMap<String, u32>,
	/// User functions list.
	functions: Cow<'static, [UserFunctionDescriptor]>,
	/// By-name functions index.
	globals_by_name: HashMap<String, u32>,
	/// User globals list.
	globals: Vec<Arc<VariableInstance>>,
}

impl<E: UserFunctionExecutor> NativeModuleInstance<E> {
	/// Create new native module
	pub fn new(base: Arc<ModuleInstanceInterface>, elements: UserDefinedElements<E>) -> Result<Self, Error> {
		if !elements.functions.is_empty() && elements.executor.is_none() {
			return Err(Error::Function("trying to construct native module with functions, but without executor".into()));
		}

		Ok(NativeModuleInstance {
			base: base,
			executor: RwLock::new(elements.executor),
			functions_by_name: elements.functions.iter().enumerate().map(|(i, f)| (f.name().to_owned(), i as u32)).collect(),
			functions: elements.functions,
			globals_by_name: elements.globals.iter().enumerate().map(|(i, (g_name, _))| (g_name.to_owned(), i as u32)).collect(),
			globals: elements.globals.into_iter().map(|(_, g)| g).collect(),
		})
	}
}

impl<E: UserFunctionExecutor> ModuleInstanceInterface for NativeModuleInstance<E> {
	fn execute_index(&self, index: u32, params: ExecutionParams) -> Result<Option<RuntimeValue>, Error> {
		self.base.execute_index(index, params)
	}

	fn execute_export(&self, name: &str, params: ExecutionParams) -> Result<Option<RuntimeValue>, Error> {
		self.base.execute_export(name, params)
	}

	fn export_entry<'b>(&self, name: &str, required_type: &ExportEntryType) -> Result<Internal, Error> {
		if let Some(index) = self.functions_by_name.get(name) {
			let composite_index = NATIVE_INDEX_FUNC_MIN + *index;
			match required_type {
				&ExportEntryType::Any => return Ok(Internal::Function(composite_index)),
				&ExportEntryType::Function(ref required_type) => {
					let actual_type = self.function_type(ItemIndex::Internal(composite_index))
						.expect(
							"by_name contains index; function_type succeeds for all functions from by_name; qed",
						);
					return if actual_type == *required_type {
						Ok(Internal::Function(composite_index))
					} else {
						Err(Error::Validation(format!(
							"Export function type {} mismatch. Expected function with signature ({:?}) -> {:?} when got with ({:?}) -> {:?}",
							index,
							required_type.params(),
							required_type.return_type(),
							actual_type.params(),
							actual_type.return_type()
						)))
					};
				}
				_ => (),
			}
		}
		if let Some(index) = self.globals_by_name.get(name) {
			let composite_index = NATIVE_INDEX_GLOBAL_MIN + *index;
			match required_type {
				&ExportEntryType::Any => {
					return Ok(Internal::Global(composite_index))
				}
				&ExportEntryType::Global(ref required_type) => {
					let actual_type = self.globals
						.get(*index as usize)
						.expect(
							"globals_by_name maps to indexes of globals; index read from globals_by_name; qed",
						)
						.variable_type();
					return if actual_type == *required_type {
						Ok(Internal::Global(composite_index))
					} else {
						Err(Error::Validation(format!(
							"Export global type {} mismatch. Expected type {:?} when got {:?}",
							index,
							required_type,
							actual_type
						)))
					};
				}
				_ => (),
			}
		}

		self.base.export_entry(name, required_type)
	}

	fn table(&self, index: ItemIndex) -> Result<Arc<TableInstance>, Error> {
		self.base.table(index)
	}

	fn memory(&self, index: ItemIndex) -> Result<Arc<MemoryInstance>, Error> {
		self.base.memory(index)
	}

	fn global<'b>(&self, global_index: ItemIndex, variable_type: Option<VariableType>, externals: Option<&'b HashMap<String, Arc<ModuleInstanceInterface + 'b>>>) -> Result<Arc<VariableInstance>, Error> {
		let index = match global_index {
			ItemIndex::IndexSpace(index) | ItemIndex::Internal(index) => index,
			ItemIndex::External(_) => unreachable!("trying to get global, exported by native module"),
		};

		if index < NATIVE_INDEX_GLOBAL_MIN {
			return self.base.global(global_index, variable_type, externals);
		}

		self.globals
			.get((index - NATIVE_INDEX_GLOBAL_MIN) as usize)
			.cloned()
			.ok_or(Error::Native(format!("trying to get native global with index {}", index)))
	}

	fn function_type(&self, function_index: ItemIndex) -> Result<FunctionSignature, Error> {
		let index = match function_index {
			ItemIndex::IndexSpace(index) | ItemIndex::Internal(index) => index,
			ItemIndex::External(_) => unreachable!("trying to call function, exported by native module"),
		};

		if index < NATIVE_INDEX_FUNC_MIN || index >= NATIVE_INDEX_GLOBAL_MIN {
			return self.base.function_type(function_index);
		}

		Ok(FunctionSignature::User(self.functions
			.get((index - NATIVE_INDEX_FUNC_MIN) as usize)
			.ok_or(Error::Native(format!("missing native function with index {}", index)))?))
	}

	fn function_type_by_index(&self, type_index: u32) -> Result<FunctionSignature, Error> {
		self.function_type(ItemIndex::Internal(type_index))
	}

	fn function_reference<'b>(&self, index: ItemIndex, externals: Option<&'b HashMap<String, Arc<ModuleInstanceInterface + 'b>>>) -> Result<InternalFunctionReference<'b>, Error> {
		self.base.function_reference(index, externals)
	}

	fn function_reference_indirect<'b>(&self, table_idx: u32, type_idx: u32, func_idx: u32, externals: Option<&'b HashMap<String, Arc<ModuleInstanceInterface + 'b>>>) -> Result<InternalFunctionReference<'b>, Error> {
		self.base.function_reference_indirect(table_idx, type_idx, func_idx, externals)
	}

	fn function_body<'b>(&'b self, _internal_index: u32) -> Result<Option<InternalFunction<'b>>, Error> {
		Ok(None)
	}

	fn call_internal_function(&self, outer: CallerContext, index: u32) -> Result<Option<RuntimeValue>, Error> {
		if index < NATIVE_INDEX_FUNC_MIN || index >= NATIVE_INDEX_GLOBAL_MIN {
			return self.base.call_internal_function(outer, index);
		}

		self.functions
			.get((index - NATIVE_INDEX_FUNC_MIN) as usize)
			.ok_or(Error::Native(format!("trying to call native function with index {}", index)).into())
			.and_then(|f| self.executor.write()
				.as_mut()
				.expect("function exists; if function exists, executor must also exists [checked in constructor]; qed")
				.execute(&f.name(), outer))
	}
}

/// Create wrapper for a module with given native user functions.
///
/// # Examples
///
/// ```rust
/// use parity_wasm::interpreter::{CallerContext, Error, ExecutionParams, ExportEntryType,
///                                FunctionSignature, ItemIndex, ModuleInstance,
///                                ModuleInstanceInterface, ProgramInstance, RuntimeValue,
///                                UserDefinedElements, UserError, UserFunctionDescriptor,
///                                UserFunctionExecutor};
///
/// struct MyExecutor;
///
/// impl UserFunctionExecutor for MyExecutor {
///     fn execute(
///         &mut self,
///         name: &str,
///         context: CallerContext,
///     ) -> Result<Option<RuntimeValue>, Error> {
///         match name {
///             "add" => {
///                 // fn add(a: u32, b: u32) -> u32
///                 let b = context.value_stack.pop_as::<u32>()? as u32;
///                 let a = context.value_stack.pop_as::<u32>()? as u32;
///                 let sum = a + b;
///                 Ok(Some(RuntimeValue::I32(sum as i32)))
///             }
///             _ => Err(Error::Trap("not implemented".into()).into()),
///         }
///     }
/// }
/// ```
pub fn native_module<'a, E: UserFunctionExecutor + 'a>(base: Arc<ModuleInstanceInterface>, user_elements: UserDefinedElements<E>) -> Result<Arc<ModuleInstanceInterface + 'a>, Error> {
	Ok(Arc::new(NativeModuleInstance::new(base, user_elements)?))
}

impl<'a> PartialEq for UserFunctionDescriptor {
	fn eq(&self, other: &Self) -> bool {
		self.params() == other.params()
			&& self.return_type() == other.return_type()
	}
}
<commit_msg>Simplified docs for interpreter::native_module() a little<commit_after>
use std::sync::Arc;
use std::collections::HashMap;
use std::borrow::Cow;
use parking_lot::RwLock;
use elements::{Internal, ValueType};
use interpreter::Error;
use interpreter::module::{ModuleInstanceInterface, ExecutionParams, ItemIndex,
	CallerContext, ExportEntryType, InternalFunctionReference, InternalFunction, FunctionSignature};
use interpreter::memory::MemoryInstance;
use interpreter::table::TableInstance;
use interpreter::value::RuntimeValue;
use interpreter::variable::{VariableInstance, VariableType};

/// Min index of native function.
pub const NATIVE_INDEX_FUNC_MIN: u32 = 10001;
/// Min index of native global.
pub const NATIVE_INDEX_GLOBAL_MIN: u32 = 20001;

/// User functions executor.
pub trait UserFunctionExecutor {
	/// Execute function with given name.
	fn execute(&mut self, name: &str, context: CallerContext) -> Result<Option<RuntimeValue>, Error>;
}

/// User function descriptor
#[derive(Debug, Clone)]
pub enum UserFunctionDescriptor {
	/// Static function definition
	Static(&'static str, &'static [ValueType], Option<ValueType>),
	/// Dynamic heap function definition
	Heap(String, Vec<ValueType>, Option<ValueType>),
}

impl UserFunctionDescriptor {
	/// New function with statically known params
	pub fn statik(name: &'static str, params: &'static [ValueType], result: Option<ValueType>) -> Self {
		UserFunctionDescriptor::Static(name, params, result)
	}

	/// New function with statically unknown params
	pub fn heap(name: String, params: Vec<ValueType>, result: Option<ValueType>) -> Self {
		UserFunctionDescriptor::Heap(name, params, result)
	}

	/// Name of the function
	pub fn name(&self) -> &str {
		match self {
			&UserFunctionDescriptor::Static(name, _, _) => name,
			&UserFunctionDescriptor::Heap(ref name, _, _) => name,
		}
	}

	/// Arguments of the function
	pub fn params(&self) -> &[ValueType] {
		match self {
			&UserFunctionDescriptor::Static(_, params, _) => params,
			&UserFunctionDescriptor::Heap(_, ref params, _) => params,
		}
	}

	/// Return type of the function
	pub fn return_type(&self) -> Option<ValueType> {
		match self {
			&UserFunctionDescriptor::Static(_, _, result) => result,
			&UserFunctionDescriptor::Heap(_, _, result) => result,
		}
	}
}

/// Set of user-defined module elements.
pub struct UserDefinedElements<E: UserFunctionExecutor> {
	/// User globals list.
	pub globals: HashMap<String, Arc<VariableInstance>>,
	/// User functions list.
	pub functions: Cow<'static, [UserFunctionDescriptor]>,
	/// Functions executor.
	pub executor: Option<E>,
}

/// Native module instance.
pub struct NativeModuleInstance<E: UserFunctionExecutor> {
	/// Underlying module reference.
	base: Arc<ModuleInstanceInterface>,
	/// User function executor.
	executor: RwLock<Option<E>>,
	/// By-name functions index.
	functions_by_name: HashMap<String, u32>,
	/// User functions list.
	functions: Cow<'static, [UserFunctionDescriptor]>,
	/// By-name functions index.
	globals_by_name: HashMap<String, u32>,
	/// User globals list.
	globals: Vec<Arc<VariableInstance>>,
}

impl<E: UserFunctionExecutor> NativeModuleInstance<E> {
	/// Create new native module
	pub fn new(base: Arc<ModuleInstanceInterface>, elements: UserDefinedElements<E>) -> Result<Self, Error> {
		if !elements.functions.is_empty() && elements.executor.is_none() {
			return Err(Error::Function("trying to construct native module with functions, but without executor".into()));
		}

		Ok(NativeModuleInstance {
			base: base,
			executor: RwLock::new(elements.executor),
			functions_by_name: elements.functions.iter().enumerate().map(|(i, f)| (f.name().to_owned(), i as u32)).collect(),
			functions: elements.functions,
			globals_by_name: elements.globals.iter().enumerate().map(|(i, (g_name, _))| (g_name.to_owned(), i as u32)).collect(),
			globals: elements.globals.into_iter().map(|(_, g)| g).collect(),
		})
	}
}

impl<E: UserFunctionExecutor> ModuleInstanceInterface for NativeModuleInstance<E> {
	fn execute_index(&self, index: u32, params: ExecutionParams) -> Result<Option<RuntimeValue>, Error> {
		self.base.execute_index(index, params)
	}

	fn execute_export(&self, name: &str, params: ExecutionParams) -> Result<Option<RuntimeValue>, Error> {
		self.base.execute_export(name, params)
	}

	fn export_entry<'b>(&self, name: &str, required_type: &ExportEntryType) -> Result<Internal, Error> {
		if let Some(index) = self.functions_by_name.get(name) {
			let composite_index = NATIVE_INDEX_FUNC_MIN + *index;
			match required_type {
				&ExportEntryType::Any => return Ok(Internal::Function(composite_index)),
				&ExportEntryType::Function(ref required_type) => {
					let actual_type = self.function_type(ItemIndex::Internal(composite_index))
						.expect(
							"by_name contains index; function_type succeeds for all functions from by_name; qed",
						);
					return if actual_type == *required_type {
						Ok(Internal::Function(composite_index))
					} else {
						Err(Error::Validation(format!(
							"Export function type {} mismatch. Expected function with signature ({:?}) -> {:?} when got with ({:?}) -> {:?}",
							index,
							required_type.params(),
							required_type.return_type(),
							actual_type.params(),
							actual_type.return_type()
						)))
					};
				}
				_ => (),
			}
		}
		if let Some(index) = self.globals_by_name.get(name) {
			let composite_index = NATIVE_INDEX_GLOBAL_MIN + *index;
			match required_type {
				&ExportEntryType::Any => {
					return Ok(Internal::Global(composite_index))
				}
				&ExportEntryType::Global(ref required_type) => {
					let actual_type = self.globals
						.get(*index as usize)
						.expect(
							"globals_by_name maps to indexes of globals; index read from globals_by_name; qed",
						)
						.variable_type();
					return if actual_type == *required_type {
						Ok(Internal::Global(composite_index))
					} else {
						Err(Error::Validation(format!(
							"Export global type {} mismatch. Expected type {:?} when got {:?}",
							index,
							required_type,
							actual_type
						)))
					};
				}
				_ => (),
			}
		}

		self.base.export_entry(name, required_type)
	}

	fn table(&self, index: ItemIndex) -> Result<Arc<TableInstance>, Error> {
		self.base.table(index)
	}

	fn memory(&self, index: ItemIndex) -> Result<Arc<MemoryInstance>, Error> {
		self.base.memory(index)
	}

	fn global<'b>(&self, global_index: ItemIndex, variable_type: Option<VariableType>, externals: Option<&'b HashMap<String, Arc<ModuleInstanceInterface + 'b>>>) -> Result<Arc<VariableInstance>, Error> {
		let index = match global_index {
			ItemIndex::IndexSpace(index) | ItemIndex::Internal(index) => index,
			ItemIndex::External(_) => unreachable!("trying to get global, exported by native module"),
		};

		if index < NATIVE_INDEX_GLOBAL_MIN {
			return self.base.global(global_index, variable_type, externals);
		}

		self.globals
			.get((index - NATIVE_INDEX_GLOBAL_MIN) as usize)
			.cloned()
			.ok_or(Error::Native(format!("trying to get native global with index {}", index)))
	}

	fn function_type(&self, function_index: ItemIndex) -> Result<FunctionSignature, Error> {
		let index = match function_index {
			ItemIndex::IndexSpace(index) | ItemIndex::Internal(index) => index,
			ItemIndex::External(_) => unreachable!("trying to call function, exported by native module"),
		};

		if index < NATIVE_INDEX_FUNC_MIN || index >= NATIVE_INDEX_GLOBAL_MIN {
			return self.base.function_type(function_index);
		}

		Ok(FunctionSignature::User(self.functions
			.get((index - NATIVE_INDEX_FUNC_MIN) as usize)
			.ok_or(Error::Native(format!("missing native function with index {}", index)))?))
	}

	fn function_type_by_index(&self, type_index: u32) -> Result<FunctionSignature, Error> {
		self.function_type(ItemIndex::Internal(type_index))
	}

	fn function_reference<'b>(&self, index: ItemIndex, externals: Option<&'b HashMap<String, Arc<ModuleInstanceInterface + 'b>>>) -> Result<InternalFunctionReference<'b>, Error> {
		self.base.function_reference(index, externals)
	}

	fn function_reference_indirect<'b>(&self, table_idx: u32, type_idx: u32, func_idx: u32, externals: Option<&'b HashMap<String, Arc<ModuleInstanceInterface + 'b>>>) -> Result<InternalFunctionReference<'b>, Error> {
		self.base.function_reference_indirect(table_idx, type_idx, func_idx, externals)
	}

	fn function_body<'b>(&'b self, _internal_index: u32) -> Result<Option<InternalFunction<'b>>, Error> {
		Ok(None)
	}

	fn call_internal_function(&self, outer: CallerContext, index: u32) -> Result<Option<RuntimeValue>, Error> {
		if index < NATIVE_INDEX_FUNC_MIN || index >= NATIVE_INDEX_GLOBAL_MIN {
			return self.base.call_internal_function(outer, index);
		}

		self.functions
			.get((index - NATIVE_INDEX_FUNC_MIN) as usize)
			.ok_or(Error::Native(format!("trying to call native function with index {}", index)).into())
			.and_then(|f| self.executor.write()
				.as_mut()
				.expect("function exists; if function exists, executor must also exists [checked in constructor]; qed")
				.execute(&f.name(), outer))
	}
}

/// Create wrapper for a module with given native user functions.
///
/// # Examples
///
/// ```rust
/// use parity_wasm::interpreter::{CallerContext, Error, RuntimeValue, UserFunctionExecutor};
///
/// struct MyExecutor;
///
/// impl UserFunctionExecutor for MyExecutor {
///     fn execute(
///         &mut self,
///         name: &str,
///         context: CallerContext,
///     ) -> Result<Option<RuntimeValue>, Error> {
///         match name {
///             "add" => {
///                 // fn add(a: u32, b: u32) -> u32
///                 let b = context.value_stack.pop_as::<u32>()?;
///                 let a = context.value_stack.pop_as::<u32>()?;
///                 let sum = a + b;
///                 Ok(Some(RuntimeValue::I32(sum as i32)))
///             }
///             _ => Err(Error::Trap("not implemented".into()).into()),
///         }
///     }
/// }
/// ```
pub fn native_module<'a, E: UserFunctionExecutor + 'a>(base: Arc<ModuleInstanceInterface>, user_elements: UserDefinedElements<E>) -> Result<Arc<ModuleInstanceInterface + 'a>, Error> {
	Ok(Arc::new(NativeModuleInstance::new(base, user_elements)?))
}

impl<'a> PartialEq for UserFunctionDescriptor {
	fn eq(&self, other: &Self) -> bool {
		self.params() == other.params()
			&& self.return_type() == other.return_type()
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 Syn Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg(feature = "extra-traits")]

extern crate syn;
use syn::*;

macro_rules! assert_let {
    ($p:pat = $e:expr) => {
        assert_let!($p = $e; {})
    };

    ($p:pat = $e:expr; $body:block) => {
        if let $p = $e
            $body
        else {
            panic!("Expected to match {} but got {:?}", stringify!($p), $e)
        }
    };
}

#[test]
#[cfg(feature = "full")]
fn test_catch_expr() {
    // Taken from tests/rust/src/test/run-pass/catch-expr.rs
    let raw = r#"
        struct catch {}

        pub fn main() {
            let catch_result = do catch {
                let x = 5;
                x
            };
            assert_eq!(catch_result, 5);

            let mut catch = true;
            while catch { catch = false; }
            assert_eq!(catch, false);

            catch = if catch { false } else { true };
            assert_eq!(catch, true);

            match catch {
                _ => {}
            };

            let catch_err = do catch {
                Err(22)?;
                Ok(1)
            };
            assert_eq!(catch_err, Err(22));

            let catch_okay: Result<i32, i32> = do catch {
                if false { Err(25)?; }
                Ok::<(), i32>(())?;
                Ok(28)
            };
            assert_eq!(catch_okay, Ok(28));

            let catch_from_loop: Result<i32, i32> = do catch {
                for i in 0..10 {
                    if i < 5 { Ok::<i32, i32>(i)?; } else { Err(i)?; }
                }
                Ok(22)
            };
            assert_eq!(catch_from_loop, Err(5));

            let cfg_init;
            let _res: Result<(), ()> = do catch {
                cfg_init = 5;
                Ok(())
            };
            assert_eq!(cfg_init, 5);

            let cfg_init_2;
            let _res: Result<(), ()> = do catch {
                cfg_init_2 = 6;
                Err(())?;
                Ok(())
            };
            assert_eq!(cfg_init_2, 6);

            let my_string = "test".to_string();
            let res: Result<&str, ()> = do catch {
                Ok(&my_string)
            };
            assert_eq!(res, Ok("test"));
        }
    "#;

    let actual: File = syn::parse_str(raw).unwrap();

    assert_let!(Item::Struct(ItemStruct { ref ident, .. }) = actual.items[0]; {
        assert_eq!(ident, "catch");
    });

    assert_let!(Item::Fn(ItemFn { ref block, .. }) = actual.items[1]; {
        assert_let!(Stmt::Local(ref local) = block.stmts[0]; {
            assert_let!(Local { init: Some((_, ref init_expr)), .. } = *local; {
                assert_let!(Expr::Catch(..) = **init_expr);
            });
        });

        assert_let!(Stmt::Local(ref local) = block.stmts[2]; {
            assert_let!(Pat::Ident(PatIdent { by_ref: None, mutability: Some(_), ref ident, .. }) = *local.pat; {
                assert_eq!(ident, "catch");
            });
        });

        assert_let!(Stmt::Expr(ref expr) = block.stmts[3]; {
            assert_let!(Expr::While(ExprWhile { ref cond, .. }) = *expr; {
                assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **cond; {
                    assert_eq!(*path, "catch".into());
                });
            });
        });

        assert_let!(Stmt::Semi(ref expr, _) = block.stmts[5]; {
            assert_let!(Expr::Assign(ExprAssign { ref left, ref right, .. }) = *expr; {
                assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **left; {
                    assert_eq!(*path, "catch".into());
                });

                assert_let!(Expr::If(ExprIf { ref cond, .. }) = **right; {
                    assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **cond; {
                        assert_eq!(*path, "catch".into());
                    });
                });
            });
        });

        assert_let!(Stmt::Semi(ref expr, _) = block.stmts[7]; {
            assert_let!(Expr::Match(ExprMatch { ref expr, .. }) = *expr; {
                assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **expr; {
                    assert_eq!(*path, "catch".into());
                });
            });
        });
    });
}
<commit_msg>Update test_expr with multiple patterns in Local<commit_after>// Copyright 2018 Syn Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg(feature = "extra-traits")]

extern crate syn;
use syn::*;

macro_rules! assert_let {
    ($p:pat = $e:expr) => {
        assert_let!($p = $e; {})
    };

    ($p:pat = $e:expr; $body:block) => {
        if let $p = $e
            $body
        else {
            panic!("Expected to match {} but got {:?}", stringify!($p), $e)
        }
    };
}

#[test]
#[cfg(feature = "full")]
fn test_catch_expr() {
    // Taken from tests/rust/src/test/run-pass/catch-expr.rs
    let raw = r#"
        struct catch {}

        pub fn main() {
            let catch_result = do catch {
                let x = 5;
                x
            };
            assert_eq!(catch_result, 5);

            let mut catch = true;
            while catch { catch = false; }
            assert_eq!(catch, false);

            catch = if catch { false } else { true };
            assert_eq!(catch, true);

            match catch {
                _ => {}
            };

            let catch_err = do catch {
                Err(22)?;
                Ok(1)
            };
            assert_eq!(catch_err, Err(22));

            let catch_okay: Result<i32, i32> = do catch {
                if false { Err(25)?; }
                Ok::<(), i32>(())?;
                Ok(28)
            };
            assert_eq!(catch_okay, Ok(28));

            let catch_from_loop: Result<i32, i32> = do catch {
                for i in 0..10 {
                    if i < 5 { Ok::<i32, i32>(i)?; } else { Err(i)?; }
                }
                Ok(22)
            };
            assert_eq!(catch_from_loop, Err(5));

            let cfg_init;
            let _res: Result<(), ()> = do catch {
                cfg_init = 5;
                Ok(())
            };
            assert_eq!(cfg_init, 5);

            let cfg_init_2;
            let _res: Result<(), ()> = do catch {
                cfg_init_2 = 6;
                Err(())?;
                Ok(())
            };
            assert_eq!(cfg_init_2, 6);

            let my_string = "test".to_string();
            let res: Result<&str, ()> = do catch {
                Ok(&my_string)
            };
            assert_eq!(res, Ok("test"));
        }
    "#;

    let actual: File = syn::parse_str(raw).unwrap();

    assert_let!(Item::Struct(ItemStruct { ref ident, .. }) = actual.items[0]; {
        assert_eq!(ident, "catch");
    });

    assert_let!(Item::Fn(ItemFn { ref block, .. }) = actual.items[1]; {
        assert_let!(Stmt::Local(ref local) = block.stmts[0]; {
            assert_let!(Local { init: Some((_, ref init_expr)), .. } = *local; {
                assert_let!(Expr::Catch(..) = **init_expr);
            });
        });

        assert_let!(Stmt::Local(ref local) = block.stmts[2]; {
            assert_let!(Pat::Ident(PatIdent { by_ref: None, mutability: Some(_), ref ident, .. }) = local.pats.iter().next().unwrap(); {
                assert_eq!(ident, "catch");
            });
        });

        assert_let!(Stmt::Expr(ref expr) = block.stmts[3]; {
            assert_let!(Expr::While(ExprWhile { ref cond, .. }) = *expr; {
                assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **cond; {
                    assert_eq!(*path, "catch".into());
                });
            });
        });

        assert_let!(Stmt::Semi(ref expr, _) = block.stmts[5]; {
            assert_let!(Expr::Assign(ExprAssign { ref left, ref right, .. }) = *expr; {
                assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **left; {
                    assert_eq!(*path, "catch".into());
                });

                assert_let!(Expr::If(ExprIf { ref cond, .. }) = **right; {
                    assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **cond; {
                        assert_eq!(*path, "catch".into());
                    });
                });
            });
        });

        assert_let!(Stmt::Semi(ref expr, _) = block.stmts[7]; {
            assert_let!(Expr::Match(ExprMatch { ref expr, .. }) = *expr; {
                assert_let!(Expr::Path(ExprPath { qself: None, ref path, .. }) = **expr; {
                    assert_eq!(*path, "catch".into());
                });
            });
        });
    });
}
<|endoftext|>


<commit_before>use std::io;
use std::io::Write;
use std::str::FromStr;
use std::time::Instant;

use anyhow::{Context, Result};
use clap::{App, Arg, ArgMatches, Values};
use futures::future::join_all;
use log::{debug, LevelFilter};

use mhost::nameserver::{self, NameServerConfig, NameServerConfigGroup};
use mhost::output::summary::SummaryOptions;
use mhost::output::{Output, OutputConfig, OutputFormat};
use mhost::resolver::{predefined, MultiQuery, ResolverConfigGroup, ResolverGroup, ResolverOpts};
use mhost::statistics::Statistics;
use mhost::{IpNetwork, RecordType};

#[tokio::main]
async fn main() -> Result<()> {
    let app = setup_clap();
    let args = app.get_matches();

    let log_level = log_level(args.occurrences_of("v"));
    setup_logging(log_level);

    debug!("Parsed args.");
    debug!("Set up logging.");

    run(args).await
}

fn log_level(verbosity: u64) -> LevelFilter {
    match verbosity {
        0 => LevelFilter::Warn,
        1 => LevelFilter::Info,
        2 => LevelFilter::Debug,
        _ => LevelFilter::Trace,
    }
}

fn setup_logging(log_level: LevelFilter) {
    let start = std::time::Instant::now();
    env_logger::Builder::from_default_env()
        .filter_module("mhost", log_level)
        .format(move |buf, rec| {
            let t = start.elapsed().as_secs_f32();
            let thread_id_string = format!("{:?}", std::thread::current().id());
            let thread_id = &thread_id_string[9..thread_id_string.len() - 1];
            writeln!(buf, "{:.03} [{:5}] ({:}) - {}", t, rec.level(), thread_id, rec.args())
        })
        .init();
}

fn setup_clap() -> App<'static, 'static> {
    App::new("mhost")
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(Arg::with_name("v")
            .short("v")
            .multiple(true)
            .help("Sets the level of verbosity"))
        .arg(Arg::with_name("domain name")
            .required_unless("list-predefined")
            .index(1)
            .value_name("NAME | IP ADDR | CIDR block")
            .help("domain name, IP address, or CIDR block")
            .long_help(
                "* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de
* IP ADDR may be any valid IPv4 or IPv4 address, e.g., 192.168.0.1
* CIDR block may be any valid IPv4 or IPv6 subnet in CIDR notation, e.g., 192.168.0.1/24
  all valid IP addresses of a CIDR block will be queried for a reverse lookup"
            )
        )
        .arg(Arg::with_name("record types")
            .short("t")
            .long("record-type")
            .value_name("RECORD TYPE")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("A,AAAA,MX")
            .help("Sets record types to lookup, will be ignored in case of IP address lookup")
        )
        .arg(Arg::with_name("no-system-resolv-opt")
            .long("no-system-resolv-opt")
            .help("Ignores options set in /etc/resolv.conf")
        )
        .arg(Arg::with_name("no-system-nameservers")
            .long("no-system-nameservers")
            .requires("system nameservers")
            .help("Ignores nameservers from /etc/resolv.conf")
        )
        .arg(Arg::with_name("system nameservers")
            .short("S")
            .long("system-nameserver")
            .value_name("IP ADDR")
            .takes_value(true)
            .multiple(true)
            .help("Adds system nameservers for system lookups; only IP addresses allowed")
        )
        .arg(Arg::with_name("nameservers")
            .short("s")
            .long("nameserver")
            .value_name("HOSTNAME | IP ADDR")
            .takes_value(true)
            .multiple(true)
            .help("Uses nameservers for lookups; if non given, only system nameservers will be used")
        )
        .arg(Arg::with_name("predefined")
            .short("p")
            .long("predefined")
            .help("Uses predefined nameservers")
        )
        .arg(Arg::with_name("list-predefined")
            .long("list-predefined")
            .conflicts_with("domain name")
            .help("Lists all predefined nameservers")
        )
        .arg(Arg::with_name("randomized-lookup")
            .short("R")
            .long("randomized-lookup")
            .help("Switches into randomize lookup mode: every query will be send just once to a one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
}

async fn run(args: ArgMatches<'_>) -> Result<()> {
    if args.is_present("list-predefined") {
        list_predefined_nameservers();
        return Ok(());
    }

    let ignore_system_resolv_opt = args.is_present("no-system-resolv-opt");
    let ignore_system_nameservers = args.is_present("no-system-nameservers");

    let domain_name = args
        .value_of("domain name")
        .context("No domain name to lookup specified")?;
    let query = if let Ok(ip_network) = IpNetwork::from_str(domain_name) {
        ptr_query(ip_network)?
    } else {
        let record_types_arg = args
            .values_of("record types")
            .context("No record types for name lookup specified")?;
        name_query(domain_name, record_types_arg)?
    };

    let system_resolver_ops = if ignore_system_resolv_opt {
        Default::default()
    } else {
        ResolverOpts::from_system_config().context("Failed to load system resolver options")?
    };
    debug!("Set system resolver opts.");

    let additional_system_nameservers: NameServerConfigGroup =
        if let Some(configs) = args.values_of("system nameservers") {
            let configs: Vec<_> = configs.map(NameServerConfig::from_str).collect();
            let configs: std::result::Result<Vec<_>, _> = configs.into_iter().collect();
            let nameservers: Vec<_> = configs.context("Failed to parse IP address for system nameserver")?;
            NameServerConfigGroup::new(nameservers)
        } else {
            NameServerConfigGroup::new(Vec::new())
        };
    debug!(
        "Prepared {} additional system nameservers.",
        additional_system_nameservers.len()
    );

    let system_nameservers: ResolverConfigGroup = if ignore_system_nameservers {
        additional_system_nameservers.into()
    } else {
        let mut system_nameservers =
            NameServerConfigGroup::from_system_config().context("Failed to load system name servers")?;
        system_nameservers.merge(additional_system_nameservers);
        debug!("Loaded {} system nameservers.", system_nameservers.len());
        system_nameservers.into()
    };

    let mut system_resolvers = ResolverGroup::from_configs(system_nameservers, system_resolver_ops, Default::default())
        .await
        .context("Failed to create system resolvers")?;
    debug!("Created {} system resolvers.", system_resolvers.len());

    let nameservers: ResolverConfigGroup = if let Some(configs) = args.values_of("nameservers") {
        let configs: Vec<_> = configs
            .map(|str| NameServerConfig::from_str_with_resolution(&system_resolvers, str))
            .collect();
        let configs: mhost::Result<Vec<_>> = join_all(configs).await.into_iter().collect();
        let nameservers: Vec<_> = configs.context("Failed to parse IP address for system nameserver")?;
        NameServerConfigGroup::new(nameservers).into()
    } else {
        NameServerConfigGroup::new(Vec::new()).into()
    };
    let resolvers = ResolverGroup::from_configs(nameservers, Default::default(), Default::default())
        .await
        .context("Failed to load resolvers")?;
    debug!("Created {} resolvers.", resolvers.len());

    system_resolvers.merge(resolvers);

    if args.is_present("predefined") {
        let configs = predefined::resolver_configs();
        let predefined = ResolverGroup::from_configs(configs, Default::default(), Default::default())
            .await
            .context("Failed to load predefined resolvers")?;
        debug!("Created {} predefined resolvers.", predefined.len());
        system_resolvers.merge(predefined);
    }

    let start_time = Instant::now();
    let lookups = if args.is_present("randomized-lookup") {
        system_resolvers.rnd_lookup(query).await
    } else {
        system_resolvers.lookup(query).await
    };
    let total_run_time = Instant::now() - start_time;

    let statistics = lookups.statistics();
    println!(
        "Received {} within {} ms of total run time.",
        statistics,
        total_run_time.as_millis()
    );

    let stdout = io::stdout();
    let mut handle = stdout.lock();

    let opts = SummaryOptions::default();
    let config = OutputConfig::summary(opts);
    let output = Output::new(config);
    output
        .output(&mut handle, &lookups)
        .context("Failed to print summary to stdout.")
}

fn list_predefined_nameservers() {
    println!("List of predefined servers:");
    for ns in nameserver::predefined::name_server_configs() {
        println!("* {}", ns);
    }
}

fn ptr_query(ip_network: IpNetwork) -> Result<MultiQuery> {
    let q = MultiQuery::multi_name(ip_network.iter(), RecordType::PTR).context("Failed to create query")?;
    debug!("Prepared query for reverse lookups.");
    Ok(q)
}

fn name_query(name: &str, record_types: Values) -> Result<MultiQuery> {
    let record_types: Vec<_> = record_types
        .map(str::to_uppercase)
        .map(|x| RecordType::from_str(&x))
        .collect();
    let record_types: std::result::Result<Vec<_>, _> = record_types.into_iter().collect();
    let record_types = record_types.context("Failed to parse record type")?;
    let record_types_len = record_types.len();

    let q = MultiQuery::multi_record(name, record_types).context("Failed to build query")?;
    debug!("Prepared query for name lookup for {} record types.", record_types_len);
    Ok(q)
}
<commit_msg>Minor CLI args tweaks<commit_after>use std::io;
use std::io::Write;
use std::str::FromStr;
use std::time::Instant;

use anyhow::{Context, Result};
use clap::{App, AppSettings, Arg, ArgMatches};
use futures::future::join_all;
use log::{debug, LevelFilter};

use mhost::{IpNetwork, RecordType};
use mhost::nameserver::{self, NameServerConfig, NameServerConfigGroup};
use mhost::output::{Output, OutputConfig, OutputFormat};
use mhost::output::summary::SummaryOptions;
use mhost::resolver::{MultiQuery, predefined, ResolverConfigGroup, ResolverGroup, ResolverOpts};
use mhost::statistics::Statistics;

static SUPPORTED_RECORD_TYPES: &[&str] = &["A", "AAAA", "ANAME", "CNAME", "MX", "NULL", "NS", "PTR", "SOA", "SRV", "TXT"];

#[tokio::main]
async fn main() -> Result<()> {
    let app = setup_clap();
    let args = app.get_matches();

    let log_level = log_level(args.occurrences_of("v"));
    setup_logging(log_level);

    debug!("Parsed args.");
    debug!("Set up logging.");

    run(args).await
}

fn log_level(verbosity: u64) -> LevelFilter {
    match verbosity {
        0 => LevelFilter::Warn,
        1 => LevelFilter::Info,
        2 => LevelFilter::Debug,
        _ => LevelFilter::Trace,
    }
}

fn setup_logging(log_level: LevelFilter) {
    let start = std::time::Instant::now();
    env_logger::Builder::from_default_env()
        .filter_module("mhost", log_level)
        .format(move |buf, rec| {
            let t = start.elapsed().as_secs_f32();
            let thread_id_string = format!("{:?}", std::thread::current().id());
            let thread_id = &thread_id_string[9..thread_id_string.len() - 1];
            writeln!(buf, "{:.03} [{:5}] ({:}) - {}", t, rec.level(), thread_id, rec.args())
        })
        .init();
}

fn setup_clap() -> App<'static, 'static> {
    App::new("mhost")
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .setting(AppSettings::ColoredHelp)
        .setting(AppSettings::DisableHelpSubcommand)
        .setting(AppSettings::DeriveDisplayOrder)
        .setting(AppSettings::UnifiedHelpMessage)
        .setting(AppSettings::VersionlessSubcommands)
        .arg(Arg::with_name("domain name")
            .required_unless("list-predefined")
            .index(1)
            .value_name("NAME | IP ADDR | CIDR BLOCK")
            .help("domain name, IP address, or CIDR block to lookup")
            .long_help(
                "* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de
* IP ADDR may be any valid IPv4 or IPv4 address, e.g., 192.168.0.1
* CIDR BLOCK may be any valid IPv4 or IPv6 subnet in CIDR notation, e.g., 192.168.0.1/24
  all valid IP addresses of a CIDR block will be queried for a reverse lookup"
            )
        )
        .arg(Arg::with_name("no-system-resolv-opt")
            .long("no-system-resolv-opt")
            .help("Ignores options set in /etc/resolv.conf")
        )
        .arg(Arg::with_name("no-system-nameservers")
            .long("no-system-nameservers")
            .requires("system nameservers")
            .help("Ignores nameservers from /etc/resolv.conf")
        )
        .arg(Arg::with_name("system nameservers")
            .short("S")
            .long("system-nameserver")
            .value_name("IP ADDR")
            .takes_value(true)
            .multiple(true)
            .help("Adds system nameserver for system lookups; only IP addresses allowed")
        )
        .arg(Arg::with_name("nameservers")
            .short("s")
            .long("nameserver")
            .value_name("HOSTNAME | IP ADDR")
            .takes_value(true)
            .multiple(true)
            .help("Adds nameserver for lookups; if non given, only system nameservers will be used")
        )
        .arg(Arg::with_name("predefined")
            .short("p")
            .long("predefined")
            .help("Adds predefined nameservers for lookups")
        )
        .arg(Arg::with_name("list-predefined")
            .long("list-predefined")
            .conflicts_with("domain name")
            .help("Lists all predefined nameservers")
        )
        .arg(Arg::with_name("record types")
            .short("t")
            .long("record-type")
            .value_name("RECORD TYPE")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("A,AAAA,MX")
            .possible_values(SUPPORTED_RECORD_TYPES)
            .help("Sets record type to lookup, will be ignored in case of IP address lookup")
        )
        .arg(Arg::with_name("all-record-types")
            .hidden(true)
            .long("all")
            .alias("xmas")
            .help("Enables lookups for all record types")
        )
        .arg(Arg::with_name("randomized-lookup")
            .short("R")
            .long("randomized-lookup")
            .help("Switches into randomize lookup mode: every query will be send just once to a one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
        .arg(Arg::with_name("v")
            .short("v")
            .multiple(true)
            .help("Sets the level of verbosity"))
}

async fn run(args: ArgMatches<'_>) -> Result<()> {
    if args.is_present("list-predefined") {
        list_predefined_nameservers();
        return Ok(());
    }

    let ignore_system_resolv_opt = args.is_present("no-system-resolv-opt");
    let ignore_system_nameservers = args.is_present("no-system-nameservers");

    let domain_name = args
        .value_of("domain name")
        .context("No domain name to lookup specified")?;
    let query = if let Ok(ip_network) = IpNetwork::from_str(domain_name) {
        ptr_query(ip_network)?
    } else {
        let record_types = if args.is_present("all-record-types") {
            SUPPORTED_RECORD_TYPES.iter().map(|x| RecordType::from_str(x).unwrap()).collect()
        } else {
            let args = args
                .values_of("record types")
                .context("No record types for name lookup specified")?;
            record_types(args)?
        };
        name_query(domain_name, record_types)?
    };

    let system_resolver_ops = if ignore_system_resolv_opt {
        Default::default()
    } else {
        ResolverOpts::from_system_config().context("Failed to load system resolver options")?
    };
    debug!("Set system resolver opts.");

    let additional_system_nameservers: NameServerConfigGroup =
        if let Some(configs) = args.values_of("system nameservers") {
            let configs: Vec<_> = configs.map(NameServerConfig::from_str).collect();
            let configs: std::result::Result<Vec<_>, _> = configs.into_iter().collect();
            let nameservers: Vec<_> = configs.context("Failed to parse IP address for system nameserver")?;
            NameServerConfigGroup::new(nameservers)
        } else {
            NameServerConfigGroup::new(Vec::new())
        };
    debug!(
        "Prepared {} additional system nameservers.",
        additional_system_nameservers.len()
    );

    let system_nameservers: ResolverConfigGroup = if ignore_system_nameservers {
        additional_system_nameservers.into()
    } else {
        let mut system_nameservers =
            NameServerConfigGroup::from_system_config().context("Failed to load system name servers")?;
        system_nameservers.merge(additional_system_nameservers);
        debug!("Loaded {} system nameservers.", system_nameservers.len());
        system_nameservers.into()
    };

    let mut system_resolvers = ResolverGroup::from_configs(system_nameservers, system_resolver_ops, Default::default())
        .await
        .context("Failed to create system resolvers")?;
    debug!("Created {} system resolvers.", system_resolvers.len());

    let nameservers: ResolverConfigGroup = if let Some(configs) = args.values_of("nameservers") {
        let configs: Vec<_> = configs
            .map(|str| NameServerConfig::from_str_with_resolution(&system_resolvers, str))
            .collect();
        let configs: mhost::Result<Vec<_>> = join_all(configs).await.into_iter().collect();
        let nameservers: Vec<_> = configs.context("Failed to parse IP address for system nameserver")?;
        NameServerConfigGroup::new(nameservers).into()
    } else {
        NameServerConfigGroup::new(Vec::new()).into()
    };
    let resolvers = ResolverGroup::from_configs(nameservers, Default::default(), Default::default())
        .await
        .context("Failed to load resolvers")?;
    debug!("Created {} resolvers.", resolvers.len());

    system_resolvers.merge(resolvers);

    if args.is_present("predefined") {
        let configs = predefined::resolver_configs();
        let predefined = ResolverGroup::from_configs(configs, Default::default(), Default::default())
            .await
            .context("Failed to load predefined resolvers")?;
        debug!("Created {} predefined resolvers.", predefined.len());
        system_resolvers.merge(predefined);
    }

    let start_time = Instant::now();
    let lookups = if args.is_present("randomized-lookup") {
        system_resolvers.rnd_lookup(query).await
    } else {
        system_resolvers.lookup(query).await
    };
    let total_run_time = Instant::now() - start_time;

    let statistics = lookups.statistics();
    println!(
        "Received {} within {} ms of total run time.",
        statistics,
        total_run_time.as_millis()
    );

    let stdout = io::stdout();
    let mut handle = stdout.lock();

    let opts = SummaryOptions::default();
    let config = OutputConfig::summary(opts);
    let output = Output::new(config);
    output
        .output(&mut handle, &lookups)
        .context("Failed to print summary to stdout.")
}

fn list_predefined_nameservers() {
    println!("List of predefined servers:");
    for ns in nameserver::predefined::name_server_configs() {
        println!("* {}", ns);
    }
}

fn ptr_query(ip_network: IpNetwork) -> Result<MultiQuery> {
    let q = MultiQuery::multi_name(ip_network.iter(), RecordType::PTR).context("Failed to create query")?;
    debug!("Prepared query for reverse lookups.");
    Ok(q)
}

fn name_query(name: &str, record_types: Vec<RecordType>) -> Result<MultiQuery> {
    let record_types_len = record_types.len();
    let q = MultiQuery::multi_record(name, record_types).context("Failed to build query")?;
    debug!("Prepared query for name lookup for {} record types.", record_types_len);
    Ok(q)
}

fn record_types<'a, I: Iterator<Item = &'a str>>(record_types: I) -> Result<Vec<RecordType>> {
    let record_types: Vec<_> = record_types
        .map(str::to_uppercase)
        .map(|x| RecordType::from_str(&x))
        .collect();
    let record_types: std::result::Result<Vec<_>, _> = record_types.into_iter().collect();
    record_types.context("Failed to parse record type")
}
<|endoftext|>


<commit_before>use lfsr::fibonacci::FibonacciLFSR;
use lfsr::galois::GaloisLFSR;

use support::eq_mod_rotation;

#[test]
fn iter_works_for_fibonacci() {
    iter_works_for!(FibonacciLFSR);
}

#[test]
fn iter_works_for_galois() {
    iter_works_for!(GaloisLFSR);
}

#[test]
fn same_connection_polynomial_produce_output_equal_mod_rotation() {
    let taps = vec![2, 3];

    let mut galois = GaloisLFSR::from_iter(taps.iter());
    let galois_output: Vec<bool> = galois.iter().take(7).collect();

    let mut fibonacci = FibonacciLFSR::from_iter(taps.iter());
    let fibonacci_output: Vec<bool> = fibonacci.iter().take(7).collect();

    assert!(eq_mod_rotation(&fibonacci_output, &galois_output));
}

#[test]
fn primitive_connection_polynomial_yields_a_maximum_sequence_for_fibonacci() {
    primitive_connection_polynomial_yields_a_maximum_sequence!(FibonacciLFSR);
}

#[test]
fn primitive_connection_polynomial_yields_a_maximum_sequence_for_galois() {
    primitive_connection_polynomial_yields_a_maximum_sequence!(GaloisLFSR);
}

#[test]
fn implements_iter_for_galois() {
    let taps = vec![2, 3];
    let mut lfsr = GaloisLFSR::from_iter(taps.iter());

    let output: Vec<bool> = lfsr.iter().take(8).collect();

    let expected = vec![true, false, false, true, false, true, true, true];

    assert!(eq_mod_rotation(&output, &expected))
}

#[test]
fn implements_into_iterator_for_galois() {
    let taps = vec![2, 3];
    let lfsr = GaloisLFSR::from_iter(taps.iter());

    let mut counter = 0;
    for _b in lfsr {
        if counter == 3 {
            break;
        }
        counter += 1;
    }
}

#[test]
fn implements_into_iterator_for_mut_refs_for_galois() {
    use lfsr::LFSR;

    let taps = vec![2, 3];
    let mut lfsr = GaloisLFSR::from_iter(taps.iter());

    let mut counter = 0;
    for _b in &mut lfsr {
        if counter == 3 {
            break;
        }
        counter += 1;
    }

    lfsr.step();
}
<commit_msg>Add FibonacciLFSR Iterator tests<commit_after>use lfsr::fibonacci::FibonacciLFSR;
use lfsr::galois::GaloisLFSR;

use support::eq_mod_rotation;

#[test]
fn iter_works_for_fibonacci() {
    iter_works_for!(FibonacciLFSR);
}

#[test]
fn iter_works_for_galois() {
    iter_works_for!(GaloisLFSR);
}

#[test]
fn same_connection_polynomial_produce_output_equal_mod_rotation() {
    let taps = vec![2, 3];

    let mut galois = GaloisLFSR::from_iter(taps.iter());
    let galois_output: Vec<bool> = galois.iter().take(7).collect();

    let mut fibonacci = FibonacciLFSR::from_iter(taps.iter());
    let fibonacci_output: Vec<bool> = fibonacci.iter().take(7).collect();

    assert!(eq_mod_rotation(&fibonacci_output, &galois_output));
}

#[test]
fn primitive_connection_polynomial_yields_a_maximum_sequence_for_fibonacci() {
    primitive_connection_polynomial_yields_a_maximum_sequence!(FibonacciLFSR);
}

#[test]
fn primitive_connection_polynomial_yields_a_maximum_sequence_for_galois() {
    primitive_connection_polynomial_yields_a_maximum_sequence!(GaloisLFSR);
}

#[test]
fn implements_iter_for_galois() {
    let taps = vec![2, 3];
    let mut lfsr = GaloisLFSR::from_iter(taps.iter());

    let output: Vec<bool> = lfsr.iter().take(8).collect();

    let expected = vec![true, false, false, true, false, true, true, true];

    assert!(eq_mod_rotation(&output, &expected))
}

#[test]
fn implements_into_iterator_for_galois() {
    let taps = vec![2, 3];
    let lfsr = GaloisLFSR::from_iter(taps.iter());

    let mut counter = 0;
    for _b in lfsr {
        if counter == 3 {
            break;
        }
        counter += 1;
    }
}

#[test]
fn implements_into_iterator_for_mut_refs_for_galois() {
    use lfsr::LFSR;

    let taps = vec![2, 3];
    let mut lfsr = GaloisLFSR::from_iter(taps.iter());

    let mut counter = 0;
    for _b in &mut lfsr {
        if counter == 3 {
            break;
        }
        counter += 1;
    }

    lfsr.step();
}

#[test]
fn implements_iter_for_fibonacci() {
    let taps = vec![2, 3];
    let mut lfsr = FibonacciLFSR::from_iter(taps.iter());

    let output: Vec<bool> = lfsr.iter().take(8).collect();

    let expected = vec![true, false, false, true, false, true, true, true];

    assert!(eq_mod_rotation(&output, &expected))
}

#[test]
fn implements_into_iterator_for_fibonacci() {
    let taps = vec![2, 3];
    let lfsr = FibonacciLFSR::from_iter(taps.iter());

    let mut counter = 0;
    for _b in lfsr {
        if counter == 3 {
            break;
        }
        counter += 1;
    }
}

#[test]
fn implements_into_iterator_for_mut_refs_for_fibonacci() {
    use lfsr::LFSR;

    let taps = vec![2, 3];
    let mut lfsr = FibonacciLFSR::from_iter(taps.iter());

    let mut counter = 0;
    for _b in &mut lfsr {
        if counter == 3 {
            break;
        }
        counter += 1;
    }

    lfsr.step();
}
<|endoftext|>