<commit_before>// The MIT License (MIT)

// Copyright (c) 2015 Rustcc Developers

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io;
use std::net::{ToSocketAddrs, SocketAddr};
use std::ops::{Deref, DerefMut};
use std::convert::From;
use std::iter::Iterator;

use mio::{self, Interest};
use mio::buf::{Buf, MutBuf, MutSliceBuf, SliceBuf};

use processor::Processor;

#[derive(Debug)]
pub struct TcpSocket(::mio::tcp::TcpSocket);

impl TcpSocket {
    /// Returns a new, unbound, non-blocking, IPv4 socket
    pub fn v4() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v4())))
    }

    /// Returns a new, unbound, non-blocking, IPv6 socket
    pub fn v6() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v6())))
    }

    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<(TcpStream, bool)> {
        super::each_addr(addr, |a| {
            match a {
                &SocketAddr::V4(..) => try!(TcpSocket::v4()).0.connect(a),
                &SocketAddr::V6(..) => try!(TcpSocket::v6()).0.connect(a),
            }
        }).map(|(stream, complete)| (TcpStream(stream), complete))
    }

    pub fn listen(self, backlog: usize) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.listen(backlog))))
    }
}

impl Deref for TcpSocket {
    type Target = ::mio::tcp::TcpSocket;

    fn deref(&self) -> &::mio::tcp::TcpSocket {
        &self.0
    }
}

impl DerefMut for TcpSocket {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpSocket {
        &mut self.0
    }
}

#[derive(Debug)]
pub struct TcpListener(::mio::tcp::TcpListener);

impl TcpListener {
    pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> {
        super::each_addr(addr, ::mio::tcp::TcpListener::bind).map(TcpListener)
    }

    pub fn accept(&self) -> io::Result<TcpStream> {
        match self.0.accept() {
            Ok(None) => {
                debug!("accept WouldBlock; going to register into eventloop");
            },
            Ok(Some(stream)) => {
                return Ok(TcpStream(stream));
            },
            Err(err) => {
                return Err(err);
            }
        }

        loop {
            try!(Processor::current().wait_event(&self.0, Interest::readable()));

            match self.0.accept() {
                Ok(None) => {
                    warn!("accept WouldBlock; Coroutine was awaked by readable event");
                },
                Ok(Some(stream)) => {
                    return Ok(TcpStream(stream));
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }
    }

    pub fn try_clone(&self) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.try_clone())))
    }

    pub fn incoming<'a>(&'a self) -> Incoming<'a> {
        Incoming(self)
    }
}

impl Deref for TcpListener {
    type Target = ::mio::tcp::TcpListener;

    fn deref(&self) -> &::mio::tcp::TcpListener {
        &self.0
    }
}

impl DerefMut for TcpListener {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpListener {
        &mut self.0
    }
}

pub struct Incoming<'a>(&'a TcpListener);

impl<'a> Iterator for Incoming<'a> {
    type Item = io::Result<TcpStream>;

    fn next(&mut self) -> Option<io::Result<TcpStream>> {
        Some(self.0.accept())
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Shutdown {
    /// Further receptions will be disallowed.
    Read,
    /// Further  transmissions will be disallowed.
    Write,
    /// Further receptions and transmissions will be disallowed.
    Both,
}

impl From<Shutdown> for mio::tcp::Shutdown {
    fn from(shutdown: Shutdown) -> mio::tcp::Shutdown {
        match shutdown {
            Shutdown::Read => mio::tcp::Shutdown::Read,
            Shutdown::Write => mio::tcp::Shutdown::Write,
            Shutdown::Both => mio::tcp::Shutdown::Both,
        }
    }
}

#[derive(Debug)]
pub struct TcpStream(mio::tcp::TcpStream);

impl TcpStream {
    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
        match TcpSocket::connect(addr) {
            Ok((stream, completed)) => {
                if !completed {
                    try!(Processor::current().wait_event(&stream.0, Interest::writable()));
                    try!(stream.take_socket_error());
                }

                Ok(stream)
            },
            Err(err) => Err(err)
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        self.0.peer_addr()
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.0.local_addr()
    }

    pub fn try_clone(&self) -> io::Result<TcpStream> {
        let stream = try!(self.0.try_clone());

        Ok(TcpStream(stream))
    }

    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.0.shutdown(From::from(how))
    }

    pub fn take_socket_error(&self) -> io::Result<()> {
        self.0.take_socket_error()
    }
}

impl io::Read for TcpStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        use mio::TryRead;

        let mut buf = MutSliceBuf::wrap(buf);
        let mut total_len = 0;
        while buf.has_remaining() {
            match self.0.try_read_buf(&mut buf) {
                Ok(None) => {
                    debug!("TcpStream read WouldBlock");
                    break;
                },
                Ok(Some(0)) => {
                    debug!("TcpStream read 0 bytes; may be EOF");
                    return Ok(total_len);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream read {} bytes", len);
                    total_len += len;
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }

        if total_len != 0 {
            // We got something, just return!
            return Ok(total_len);
        }

        debug!("Read: Going to register event");
        try!(Processor::current().wait_event(&self.0, Interest::readable()));
        debug!("Read: Got read event");

        while buf.has_remaining() {
            match self.0.try_read_buf(&mut buf) {
                Ok(None) => {
                    debug!("TcpStream read WouldBlock");
                    break;
                },
                Ok(Some(0)) => {
                    debug!("TcpStream read 0 bytes; may be EOF");
                    break;
                },
                Ok(Some(len)) => {
                    debug!("TcpStream read {} bytes", len);
                    total_len += len;
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }

        Ok(total_len)
    }
}

impl io::Write for TcpStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        use mio::TryWrite;

        let mut buf = SliceBuf::wrap(buf);
        let mut total_len = 0;

        while buf.has_remaining() {
            match self.0.try_write_buf(&mut buf) {
                Ok(None) => {
                    debug!("TcpStream write WouldBlock");
                    break;
                },
                Ok(Some(0)) => {
                    debug!("TcpStream write 0 bytes; may be EOF");
                    return Ok(total_len);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream written {} bytes", len);
                    total_len += len;
                },
                Err(err) => {
                    return Err(err)
                }
            }
        }

        if total_len != 0 {
            // We have written something, return it!
            return Ok(total_len)
        }

        debug!("Write: Going to register event");
        try!(Processor::current().wait_event(&self.0, Interest::writable()));
        debug!("Write: Got write event");

        while buf.has_remaining() {
            match self.0.try_write_buf(&mut buf) {
                Ok(None) => {
                    debug!("TcpStream write WouldBlock");
                    break;
                },
                Ok(Some(0)) => {
                    debug!("TcpStream write 0 bytes; may be EOF");
                    break;
                },
                Ok(Some(len)) => {
                    debug!("TcpStream written {} bytes", len);
                    total_len += len;
                },
                Err(err) => {
                    return Err(err)
                }
            }
        }

        Ok(total_len)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Deref for TcpStream {
    type Target = ::mio::tcp::TcpStream;

    fn deref(&self) -> &::mio::tcp::TcpStream {
        &self.0
    }
}

impl DerefMut for TcpStream {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpStream {
        &mut self.0
    }
}
<commit_msg>fixed bug in tcp<commit_after>// The MIT License (MIT)

// Copyright (c) 2015 Rustcc Developers

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io;
use std::net::{ToSocketAddrs, SocketAddr};
use std::ops::{Deref, DerefMut};
use std::convert::From;
use std::iter::Iterator;

use mio::{self, Interest};

use processor::Processor;

#[derive(Debug)]
pub struct TcpSocket(::mio::tcp::TcpSocket);

impl TcpSocket {
    /// Returns a new, unbound, non-blocking, IPv4 socket
    pub fn v4() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v4())))
    }

    /// Returns a new, unbound, non-blocking, IPv6 socket
    pub fn v6() -> io::Result<TcpSocket> {
        Ok(TcpSocket(try!(::mio::tcp::TcpSocket::v6())))
    }

    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<(TcpStream, bool)> {
        super::each_addr(addr, |a| {
            match a {
                &SocketAddr::V4(..) => try!(TcpSocket::v4()).0.connect(a),
                &SocketAddr::V6(..) => try!(TcpSocket::v6()).0.connect(a),
            }
        }).map(|(stream, complete)| (TcpStream(stream), complete))
    }

    pub fn listen(self, backlog: usize) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.listen(backlog))))
    }
}

impl Deref for TcpSocket {
    type Target = ::mio::tcp::TcpSocket;

    fn deref(&self) -> &::mio::tcp::TcpSocket {
        &self.0
    }
}

impl DerefMut for TcpSocket {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpSocket {
        &mut self.0
    }
}

#[derive(Debug)]
pub struct TcpListener(::mio::tcp::TcpListener);

impl TcpListener {
    pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> {
        super::each_addr(addr, ::mio::tcp::TcpListener::bind).map(TcpListener)
    }

    pub fn accept(&self) -> io::Result<TcpStream> {
        match self.0.accept() {
            Ok(None) => {
                debug!("accept WouldBlock; going to register into eventloop");
            },
            Ok(Some(stream)) => {
                return Ok(TcpStream(stream));
            },
            Err(err) => {
                return Err(err);
            }
        }

        loop {
            try!(Processor::current().wait_event(&self.0, Interest::readable()));

            match self.0.accept() {
                Ok(None) => {
                    warn!("accept WouldBlock; Coroutine was awaked by readable event");
                },
                Ok(Some(stream)) => {
                    return Ok(TcpStream(stream));
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }
    }

    pub fn try_clone(&self) -> io::Result<TcpListener> {
        Ok(TcpListener(try!(self.0.try_clone())))
    }

    pub fn incoming<'a>(&'a self) -> Incoming<'a> {
        Incoming(self)
    }
}

impl Deref for TcpListener {
    type Target = ::mio::tcp::TcpListener;

    fn deref(&self) -> &::mio::tcp::TcpListener {
        &self.0
    }
}

impl DerefMut for TcpListener {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpListener {
        &mut self.0
    }
}

pub struct Incoming<'a>(&'a TcpListener);

impl<'a> Iterator for Incoming<'a> {
    type Item = io::Result<TcpStream>;

    fn next(&mut self) -> Option<io::Result<TcpStream>> {
        Some(self.0.accept())
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Shutdown {
    /// Further receptions will be disallowed.
    Read,
    /// Further  transmissions will be disallowed.
    Write,
    /// Further receptions and transmissions will be disallowed.
    Both,
}

impl From<Shutdown> for mio::tcp::Shutdown {
    fn from(shutdown: Shutdown) -> mio::tcp::Shutdown {
        match shutdown {
            Shutdown::Read => mio::tcp::Shutdown::Read,
            Shutdown::Write => mio::tcp::Shutdown::Write,
            Shutdown::Both => mio::tcp::Shutdown::Both,
        }
    }
}

#[derive(Debug)]
pub struct TcpStream(mio::tcp::TcpStream);

impl TcpStream {
    pub fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
        match TcpSocket::connect(addr) {
            Ok((stream, completed)) => {
                if !completed {
                    try!(Processor::current().wait_event(&stream.0, Interest::writable()));
                    try!(stream.take_socket_error());
                }

                Ok(stream)
            },
            Err(err) => Err(err)
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        self.0.peer_addr()
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.0.local_addr()
    }

    pub fn try_clone(&self) -> io::Result<TcpStream> {
        let stream = try!(self.0.try_clone());

        Ok(TcpStream(stream))
    }

    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.0.shutdown(From::from(how))
    }

    pub fn take_socket_error(&self) -> io::Result<()> {
        self.0.take_socket_error()
    }
}

impl io::Read for TcpStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        use mio::TryRead;

        match self.0.try_read(buf) {
            Ok(None) => {
                debug!("TcpStream read WouldBlock");
            },
            Ok(Some(0)) => {
                debug!("TcpStream read 0 bytes; may be EOF");
                return Ok(0);
            },
            Ok(Some(len)) => {
                debug!("TcpStream read {} bytes", len);
                return Ok(len);
            },
            Err(err) => {
                return Err(err);
            }
        }

        loop {
            debug!("Read: Going to register event");
            try!(Processor::current().wait_event(&self.0, Interest::readable()));
            debug!("Read: Got read event");

            match self.0.try_read(buf) {
                Ok(None) => {
                    warn!("TcpStream read WouldBlock");
                },
                Ok(Some(0)) => {
                    debug!("TcpStream read 0 bytes; may be EOF");
                    return Ok(0);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream read {} bytes", len);
                    return Ok(len);
                },
                Err(err) => {
                    return Err(err);
                }
            }
        }
    }
}

impl io::Write for TcpStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        use mio::TryWrite;

        match self.0.try_write(buf) {
            Ok(None) => {
                debug!("TcpStream write WouldBlock");
            },
            Ok(Some(0)) => {
                debug!("TcpStream write 0 bytes; may be EOF");
                return Ok(0);
            },
            Ok(Some(len)) => {
                debug!("TcpStream written {} bytes", len);
                return Ok(len);
            },
            Err(err) => {
                return Err(err)
            }
        }

        loop {
             debug!("Write: Going to register event");
            try!(Processor::current().wait_event(&self.0, Interest::writable()));
            debug!("Write: Got write event");

            match self.0.try_write(buf) {
                Ok(None) => {
                    warn!("TcpStream write WouldBlock");
                },
                Ok(Some(0)) => {
                    debug!("TcpStream write 0 bytes; may be EOF");
                    return Ok(0);
                },
                Ok(Some(len)) => {
                    debug!("TcpStream written {} bytes", len);
                    return Ok(len);
                },
                Err(err) => {
                    return Err(err)
                }
            }
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Deref for TcpStream {
    type Target = ::mio::tcp::TcpStream;

    fn deref(&self) -> &::mio::tcp::TcpStream {
        &self.0
    }
}

impl DerefMut for TcpStream {
    fn deref_mut(&mut self) -> &mut ::mio::tcp::TcpStream {
        &mut self.0
    }
}
<|endoftext|>


<commit_before>use lookup;
use txt_records::{Spf, Word, Mechanism, Modifier};
use statistics::Statistics;

use ansi_term::Colour;
use chrono::{Local, Duration};
use chrono_humanize::HumanTime;
use error_chain::ChainedError;
use itertools::Itertools;
use std::cmp::Ordering;
use std::io::{self, Write};
use tabwriter::TabWriter;
use trust_dns::rr::{RData, Record};

pub trait OutputModule {
    fn output(&self, w: &mut Write) -> Result<()>;
    fn is_human(&self) -> bool;
}

pub struct DetailsOutput<'a> {
    responses: &'a [lookup::Result<lookup::Response>],
    human: bool,
}

impl<'a> DetailsOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        DetailsOutput { responses, human }
    }
}

impl<'a> OutputModule for DetailsOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        for response in self.responses {
            match *response {
                Ok(ref r) => write_response(w, r, self.human).chain_err(|| ErrorKind::OutputError)?,
                Err(ref e) => print_error(w, e)?,
            }
        }

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

pub struct SummaryOutput<'a> {
    statistics: Statistics<'a>,
    human: bool,
}

impl<'a> SummaryOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        let statistics = Statistics::from(responses);
        SummaryOutput { statistics, human }
    }
}

impl<'a> OutputModule for SummaryOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        writeln!(w, "Received {} (min {}, max {} records) answers from {} servers",
                 self.statistics.num_of_ok_samples,
                 self.statistics.min_num_of_records,
                 self.statistics.max_num_of_records,
                 self.statistics.num_of_samples,
        ).chain_err(|| ErrorKind::OutputError)?;
        let records: Vec<_> = self.statistics
            .record_counts
            .values()
            // TODO: Why do I need to specify a closure and not just a function?
            .sorted_by(|a, b| compare_records(a.0, b.0))
            .iter()
            .map(|rr| {
                let record = rr.0;
                let count = rr.1;
                format!("* {} ({})", fmt_record(record, self.human), count)
            })
            .collect();

        let mut tw = TabWriter::new(vec![]).padding(1);
        write!(&mut tw, "{}", records.join("\n")).chain_err(|| ErrorKind::OutputError)?;
        let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

        writeln!(w, "{}", out_str).chain_err(|| ErrorKind::OutputError)?;

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

fn write_response(f: &mut Write, r: &lookup::Response, human: bool) -> io::Result<()> {
    if r.answers.is_empty() {
        return writeln!(f, "DNS server {} has no records.", r.server);
    }
    let _ = write!(f, "DNS server {} responded with\n", r.server);
    let answers: Vec<String> = r.answers
        .iter()
        .sorted_by(|a, b| compare_records(a, b))
        .iter()
        .map(|answer|
            format!("* {} [expires {}]",
                    fmt_record(answer, human),
                    if human {
                        humanize_ttl(answer.ttl() as i64)
                    } else {
                        format!("in {} sec", answer.ttl())
                    }
            )
        )
        .collect();

    let mut tw = TabWriter::new(vec![]).padding(1);
    let _ = write!(&mut tw, "{}", answers.join("\n"));
    let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

    writeln!(f, "{}", out_str)
}

fn compare_records(a: &Record, b: &Record) -> Ordering {
    let a = record_type_to_ordinal(a);
    let b = record_type_to_ordinal(b);

    a.cmp(&b)
}

fn record_type_to_ordinal(r: &Record) -> u16 {
    match *r.rdata() {
        RData::SOA(_) => 1000,
        RData::NS(_) => 2000,
        RData::MX(ref mx) => 3000 + mx.preference(),
        RData::TXT(_) => 4000,
        RData::CNAME(_) => 5000,
        RData::A(_) => 6000,
        RData::AAAA(_) => 7000,
        RData::PTR(_) => 8000,
        _ => ::std::u16::MAX,
    }
}

fn humanize_ttl(ttl: i64) -> String {
    let dt = Local::now() + Duration::seconds(ttl);
    let ht = HumanTime::from(dt);

    format!("{}", ht)
}

fn fmt_record(r: &Record, human: bool) -> String {
    match *r.rdata() {
        RData::A(ip) => format!("IPv4:\t{}", ip),
        RData::AAAA(ip) => format!("IPv6:\t{}", ip),
        RData::CNAME(ref name) => {
            format!("CNAME:\t{}", Colour::Blue.paint(format!("{}", name)))
        }
        RData::MX(ref mx) => {
            format!(
                "MX:\t{} with preference {}",
                Colour::Yellow.paint(format!("{}", mx.exchange())),
                Colour::Yellow.paint(format!("{}", mx.preference()))
            )
        }
        RData::NS(ref name) => format!("NS:\t{}", Colour::Cyan.paint(format!("{}", name))),
        RData::SOA(ref soa) => {
            if human {
                format!(
                    "SOA:\torigin NS {}, responsible party {}, serial {}, refresh {}, retry {}, expire {}, min {}",
                    Colour::Green.paint(format!("{}", soa.mname())),
                    Colour::Green.paint(format!("{}", soa.rname())),
                    Colour::Green.paint(format!("{}", soa.serial())),
                    Colour::Green.paint(humanize_ttl(soa.refresh() as i64)),
                    Colour::Green.paint(humanize_ttl(soa.retry() as i64)),
                    Colour::Green.paint(humanize_ttl(soa.expire() as i64)),
                    Colour::Green.paint(humanize_ttl(soa.minimum() as i64))
                )
            } else {
                format!(
                    "SOA:\torigin NS {}, responsible party {}, serial {}, refresh {} sec, retry {} sec, expire {} sec, min {} sec",
                    Colour::Green.paint(format!("{}", soa.mname())),
                    Colour::Green.paint(format!("{}", soa.rname())),
                    Colour::Green.paint(format!("{}", soa.serial())),
                    Colour::Green.paint(format!("{}", soa.refresh())),
                    Colour::Green.paint(format!("{}", soa.retry())),
                    Colour::Green.paint(format!("{}", soa.expire())),
                    Colour::Green.paint(format!("{}", soa.minimum()))
                )

            }
        }
        RData::TXT(ref txt) => {
            format!("TXT:\t{}", Colour::Purple.paint(
                if human {
                    fmt_txt(txt.txt_data())
                } else {
                    txt.txt_data().join(" ")
                }
            ))
        }
        RData::PTR(ref ptr) => format!("PTR:\t{}", ptr.to_string()),
        ref x => {
            format!(
                " * unclassified answer: {}",
                Colour::Red.paint(format!("{:?}", x))
            )
        }
    }
}

fn fmt_txt(txts: &[String]) -> String {
    let fmts: Vec<_> = txts
        .iter()
        .map(|txt| {
            if let Ok(spf) = Spf::from_str(txt) {
                fmt_txt_spf(&spf)
            } else {
                txt.to_string()
            }
        })
        .collect();
    fmts.iter().join("\t* ")
}

fn fmt_txt_spf(spf: &Spf) -> String {
    let words: Vec<_> = spf.words
        .iter()
        .map(|w| {
            match *w {
                Word::Word(ref q, Mechanism::All) => format!("{:?} for all", q),
                Word::Word(ref q, Mechanism::A) => format!("{:?} for A/AAAA record", q),
                Word::Word(ref q, Mechanism::IPv4(range)) if range.contains('/') => format!("{:?} for IPv4 range {}", q, range),
                Word::Word(ref q, Mechanism::IPv4(range)) => format!("{:?} for IPv4 {}", q, range),
                Word::Word(ref q, Mechanism::IPv6(range)) if range.contains('/') => format!("{:?} for IPv6 range {}", q, range),
                Word::Word(ref q, Mechanism::IPv6(range)) => format!("{:?} for IPv6 {}", q, range),
                Word::Word(ref q, Mechanism::MX) => format!("{:?} for mail exchanges", q),
                Word::Word(ref q, Mechanism::PTR) => format!("{:?} for reverse mapping", q),
                Word::Word(ref q, Mechanism::Exists(domain)) => format!("{:?} for A/AAAA record according to {}", q, domain),
                Word::Word(ref q, Mechanism::Include(domain)) => format!("{:?} for include from {}", q, domain),
                Word::Modifier(Modifier::Redirect(query)) => format!("redirect to query {}", query),
                Word::Modifier(Modifier::Exp(explanation)) => format!("explanation according to {}", explanation),
            }
        })
        .collect();
    format!("SPF version: {}\n\t* {}", spf.version, words.join("\n\t* "))
}

pub fn print_error<T: ChainedError>(w: &mut Write, err: &T) -> Result<()> {
    write!(w, "{} ", err).chain_err(|| ErrorKind::OutputError)?;
    for e in err.iter().skip(1) {
        write!(w, "because {}", e).chain_err(|| ErrorKind::OutputError)?;
    }
    writeln!(w).chain_err(|| ErrorKind::OutputError)?;

    // The backtrace is only available if run with `RUST_BACKTRACE=1`.
    if let Some(backtrace) = err.backtrace() {
        writeln!(w, "backtrace: {:?}", backtrace).chain_err(|| ErrorKind::OutputError)?;
    }

    Ok(())
}

error_chain! {
    errors {
        OutputError {
            description("Failed to write output")
            display("Failed to write output")
        }
    }
}
<commit_msg>Minor output fix for unsupported RR<commit_after>use lookup;
use txt_records::{Spf, Word, Mechanism, Modifier};
use statistics::Statistics;

use ansi_term::Colour;
use chrono::{Local, Duration};
use chrono_humanize::HumanTime;
use error_chain::ChainedError;
use itertools::Itertools;
use std::cmp::Ordering;
use std::io::{self, Write};
use tabwriter::TabWriter;
use trust_dns::rr::{RData, Record};

pub trait OutputModule {
    fn output(&self, w: &mut Write) -> Result<()>;
    fn is_human(&self) -> bool;
}

pub struct DetailsOutput<'a> {
    responses: &'a [lookup::Result<lookup::Response>],
    human: bool,
}

impl<'a> DetailsOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        DetailsOutput { responses, human }
    }
}

impl<'a> OutputModule for DetailsOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        for response in self.responses {
            match *response {
                Ok(ref r) => write_response(w, r, self.human).chain_err(|| ErrorKind::OutputError)?,
                Err(ref e) => print_error(w, e)?,
            }
        }

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

pub struct SummaryOutput<'a> {
    statistics: Statistics<'a>,
    human: bool,
}

impl<'a> SummaryOutput<'a> {
    pub fn new(responses: &'a [lookup::Result<lookup::Response>], human: bool) -> Self {
        let statistics = Statistics::from(responses);
        SummaryOutput { statistics, human }
    }
}

impl<'a> OutputModule for SummaryOutput<'a> {
    fn output(&self, w: &mut Write) -> Result<()> {
        writeln!(w, "Received {} (min {}, max {} records) answers from {} servers",
                 self.statistics.num_of_ok_samples,
                 self.statistics.min_num_of_records,
                 self.statistics.max_num_of_records,
                 self.statistics.num_of_samples,
        ).chain_err(|| ErrorKind::OutputError)?;
        let records: Vec<_> = self.statistics
            .record_counts
            .values()
            // TODO: Why do I need to specify a closure and not just a function?
            .sorted_by(|a, b| compare_records(a.0, b.0))
            .iter()
            .map(|rr| {
                let record = rr.0;
                let count = rr.1;
                format!("* {} ({})", fmt_record(record, self.human), count)
            })
            .collect();

        let mut tw = TabWriter::new(vec![]).padding(1);
        write!(&mut tw, "{}", records.join("\n")).chain_err(|| ErrorKind::OutputError)?;
        let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

        writeln!(w, "{}", out_str).chain_err(|| ErrorKind::OutputError)?;

        Ok(())
    }

    fn is_human(&self) -> bool {
        self.human
    }
}

fn write_response(f: &mut Write, r: &lookup::Response, human: bool) -> io::Result<()> {
    if r.answers.is_empty() {
        return writeln!(f, "DNS server {} has no records.", r.server);
    }
    let _ = write!(f, "DNS server {} responded with\n", r.server);
    let answers: Vec<String> = r.answers
        .iter()
        .sorted_by(|a, b| compare_records(a, b))
        .iter()
        .map(|answer|
            format!("* {} [expires {}]",
                    fmt_record(answer, human),
                    if human {
                        humanize_ttl(answer.ttl() as i64)
                    } else {
                        format!("in {} sec", answer.ttl())
                    }
            )
        )
        .collect();

    let mut tw = TabWriter::new(vec![]).padding(1);
    let _ = write!(&mut tw, "{}", answers.join("\n"));
    let out_str = String::from_utf8(tw.into_inner().unwrap()).unwrap();

    writeln!(f, "{}", out_str)
}

fn compare_records(a: &Record, b: &Record) -> Ordering {
    let a = record_type_to_ordinal(a);
    let b = record_type_to_ordinal(b);

    a.cmp(&b)
}

fn record_type_to_ordinal(r: &Record) -> u16 {
    match *r.rdata() {
        RData::SOA(_) => 1000,
        RData::NS(_) => 2000,
        RData::MX(ref mx) => 3000 + mx.preference(),
        RData::TXT(_) => 4000,
        RData::CNAME(_) => 5000,
        RData::A(_) => 6000,
        RData::AAAA(_) => 7000,
        RData::PTR(_) => 8000,
        _ => ::std::u16::MAX,
    }
}

fn humanize_ttl(ttl: i64) -> String {
    let dt = Local::now() + Duration::seconds(ttl);
    let ht = HumanTime::from(dt);

    format!("{}", ht)
}

fn fmt_record(r: &Record, human: bool) -> String {
    match *r.rdata() {
        RData::A(ip) => format!("IPv4:\t{}", ip),
        RData::AAAA(ip) => format!("IPv6:\t{}", ip),
        RData::CNAME(ref name) => {
            format!("CNAME:\t{}", Colour::Blue.paint(format!("{}", name)))
        }
        RData::MX(ref mx) => {
            format!(
                "MX:\t{} with preference {}",
                Colour::Yellow.paint(format!("{}", mx.exchange())),
                Colour::Yellow.paint(format!("{}", mx.preference()))
            )
        }
        RData::NS(ref name) => format!("NS:\t{}", Colour::Cyan.paint(format!("{}", name))),
        RData::SOA(ref soa) => {
            if human {
                format!(
                    "SOA:\torigin NS {}, responsible party {}, serial {}, refresh {}, retry {}, expire {}, min {}",
                    Colour::Green.paint(format!("{}", soa.mname())),
                    Colour::Green.paint(format!("{}", soa.rname())),
                    Colour::Green.paint(format!("{}", soa.serial())),
                    Colour::Green.paint(humanize_ttl(soa.refresh() as i64)),
                    Colour::Green.paint(humanize_ttl(soa.retry() as i64)),
                    Colour::Green.paint(humanize_ttl(soa.expire() as i64)),
                    Colour::Green.paint(humanize_ttl(soa.minimum() as i64))
                )
            } else {
                format!(
                    "SOA:\torigin NS {}, responsible party {}, serial {}, refresh {} sec, retry {} sec, expire {} sec, min {} sec",
                    Colour::Green.paint(format!("{}", soa.mname())),
                    Colour::Green.paint(format!("{}", soa.rname())),
                    Colour::Green.paint(format!("{}", soa.serial())),
                    Colour::Green.paint(format!("{}", soa.refresh())),
                    Colour::Green.paint(format!("{}", soa.retry())),
                    Colour::Green.paint(format!("{}", soa.expire())),
                    Colour::Green.paint(format!("{}", soa.minimum()))
                )

            }
        }
        RData::TXT(ref txt) => {
            format!("TXT:\t{}", Colour::Purple.paint(
                if human {
                    fmt_txt(txt.txt_data())
                } else {
                    txt.txt_data().join(" ")
                }
            ))
        }
        RData::PTR(ref ptr) => format!("PTR:\t{}", ptr.to_string()),
        ref x => {
            format!(
                "Unsupported RR:\t{}",
                Colour::Red.paint(format!("{:?}", x))
            )
        }
    }
}

fn fmt_txt(txts: &[String]) -> String {
    let fmts: Vec<_> = txts
        .iter()
        .map(|txt| {
            if let Ok(spf) = Spf::from_str(txt) {
                fmt_txt_spf(&spf)
            } else {
                txt.to_string()
            }
        })
        .collect();
    fmts.iter().join("\t* ")
}

fn fmt_txt_spf(spf: &Spf) -> String {
    let words: Vec<_> = spf.words
        .iter()
        .map(|w| {
            match *w {
                Word::Word(ref q, Mechanism::All) => format!("{:?} for all", q),
                Word::Word(ref q, Mechanism::A) => format!("{:?} for A/AAAA record", q),
                Word::Word(ref q, Mechanism::IPv4(range)) if range.contains('/') => format!("{:?} for IPv4 range {}", q, range),
                Word::Word(ref q, Mechanism::IPv4(range)) => format!("{:?} for IPv4 {}", q, range),
                Word::Word(ref q, Mechanism::IPv6(range)) if range.contains('/') => format!("{:?} for IPv6 range {}", q, range),
                Word::Word(ref q, Mechanism::IPv6(range)) => format!("{:?} for IPv6 {}", q, range),
                Word::Word(ref q, Mechanism::MX) => format!("{:?} for mail exchanges", q),
                Word::Word(ref q, Mechanism::PTR) => format!("{:?} for reverse mapping", q),
                Word::Word(ref q, Mechanism::Exists(domain)) => format!("{:?} for A/AAAA record according to {}", q, domain),
                Word::Word(ref q, Mechanism::Include(domain)) => format!("{:?} for include from {}", q, domain),
                Word::Modifier(Modifier::Redirect(query)) => format!("redirect to query {}", query),
                Word::Modifier(Modifier::Exp(explanation)) => format!("explanation according to {}", explanation),
            }
        })
        .collect();
    format!("SPF version: {}\n\t* {}", spf.version, words.join("\n\t* "))
}

pub fn print_error<T: ChainedError>(w: &mut Write, err: &T) -> Result<()> {
    write!(w, "{} ", err).chain_err(|| ErrorKind::OutputError)?;
    for e in err.iter().skip(1) {
        write!(w, "because {}", e).chain_err(|| ErrorKind::OutputError)?;
    }
    writeln!(w).chain_err(|| ErrorKind::OutputError)?;

    // The backtrace is only available if run with `RUST_BACKTRACE=1`.
    if let Some(backtrace) = err.backtrace() {
        writeln!(w, "backtrace: {:?}", backtrace).chain_err(|| ErrorKind::OutputError)?;
    }

    Ok(())
}

error_chain! {
    errors {
        OutputError {
            description("Failed to write output")
            display("Failed to write output")
        }
    }
}
<|endoftext|>


<commit_before>//! Executes compiled JMESPath bytecode instructions.

extern crate rustc_serialize;

use self::rustc_serialize::json::Json;

use ast::{Ast, Comparator, KeyValuePair};

#[derive(Clone, PartialEq, Debug)]
pub enum Opcode {
    // Pops N operands from stack and calls a built-in function by name.
    Call(u8, String),
    // Pops two elements from the stack and pushes true/false.
    Cmp(Comparator),
    // Duplicates the top of the stack
    Dup,
    // Pops TOS and iterates over each element if it is an array. If it
    // is not an array, then jumps to the given address.
    Each(usize),
    // Pops TOS and pushes the <operand> key value of TOS onto the stack.
    // If TOS is not an "object", null is pushed onto the stack.
    Field(String),
    // Pops TOS and pushes the <operand> index value of TOS onto the stack.
    // If TOS is not an "array", null is pushed onto the stack.
    Index(usize),
    // Pops TOS and pushes the <operand> index from the end of the length of
    // of TOS onto the stack. If TOS is not an "array", null is pushed onto
    /// the stack.
    NegativeIndex(usize),
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "array". Pushes T1 onto T2 and pushes T2 back onto the stack.
    InjectIdx(usize),
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "object". Adds T1 at the given operand key onto T2 and pushes T2
    // back onto the stack.
    InjectKey(String),
    // Unconditional jump to an address.
    Jump(usize),
    // Pops the TOS and jumps if 0 or false
    Jumpz(usize),
    // Pops the call stack
    Ret(usize),
    // Pops the TOS
    Pop,
    // Pushes a value onto the stack
    Push(Json),
    // Pops TOS and pushes true/false if the value is truthy.
    Truthy,
    // Pops TOS and pushes the type of the popped value.
    Type
}

pub type Program = Vec<Opcode>;

/// JMESPath VM
pub struct Vm<'a> {
    /// The compiled program.
    opcodes: &'a Program,
    /// The opcode index.
    index: usize,
    /// The VM stack consisting of JSON typed values.
    stack: Vec<Json>
}

impl<'a> Vm<'a> {
    /// Creates a new VM using the given program and data.
    pub fn new(opcodes: &'a Vec<Opcode>, data: Json) -> Vm<'a> {
        Vm {
            opcodes: opcodes,
            stack: vec![data],
            index: 0
        }
    }

    pub fn run(&mut self) -> Result<Json, String> {
        macro_rules! tos {
            () => {{
                self.stack.pop().unwrap_or(Json::Null)
            }};
        }
        while self.index < self.opcodes.len() {
            match &self.opcodes[self.index] {
                &Opcode::Push(ref j) => {
                    self.stack.push(j.clone());
                },
                &Opcode::Pop => {
                    self.stack.pop();
                },
                &Opcode::Dup => {
                    if self.stack.is_empty() {
                        self.stack.push(Json::Null);
                    } else {
                        let tos = self.stack.pop().unwrap();
                        self.stack.push(tos.clone());
                        self.stack.push(tos);
                    }
                },
                &Opcode::Field(ref f) => {
                    match tos!().find(f) {
                        Some(j) => self.stack.push(j.clone()),
                        _ => self.stack.push(Json::Null)
                    };
                },
                &Opcode::Index(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|_| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::NegativeIndex(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|i| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::Truthy => {
                    let tos = tos!();
                    self.stack.push(Json::Boolean(match tos {
                        Json::Boolean(b) if b == false => false,
                        Json::String(ref s) if s.len() == 0 => false,
                        Json::Array(ref a) if a.is_empty() => false,
                        Json::Object(ref o) if o.is_empty() => false,
                        Json::Null => false,
                        _ => true
                    }));
                },
                &Opcode::Type => {
                    let tos = tos!();
                    self.stack.push(match tos {
                        Json::Boolean(_) => Json::String("boolean".to_string()),
                        Json::String(_) => Json::String("string".to_string()),
                        Json::I64(_) | Json::U64(_) | Json::F64(_) => {
                            Json::String("number".to_string())
                        },
                        Json::Array(_) => Json::String("array".to_string()),
                        Json::Object(_) => Json::String("object".to_string()),
                        Json::Null => Json::String("null".to_string()),
                    });
                },
                &Opcode::Jump(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));

                    }
                    self.index = address;
                    continue;
                },
                &Opcode::Jumpz(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));
                    }
                    let tos = tos!();
                    self.index = match tos {
                        Json::Boolean(b) if b == false => address,
                        Json::I64(d) if d == 0 => address,
                        Json::F64(f) if f == 0f64 => address,
                        Json::U64(u) if u == 0 => address,
                        _ => self.index + 1
                    };
                    continue;
                },
                _ => panic!("Not implemented yet!")
            }
            self.index += 1;
        }

        Ok(self.stack.pop().unwrap_or(Json::Null))
    }
}

#[cfg(test)]
mod test {
    extern crate rustc_serialize;
    use super::{Opcode, Vm};
    use self::rustc_serialize::json::Json;

    #[test] fn truthy_test() {
        let opcodes = vec![Opcode::Truthy];
        let tests = vec![("\"foo\"", true), ("\"\"", false),
                         ("[1]", true), ("[]", false),
                         ("{\"foo\": 1}", true), ("{}", false),
                         ("true", true), ("false", false),
                         ("null", false), ("1", true)];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, Json::from_str(js).unwrap());
            assert_eq!(Json::Boolean(result), vm.run().unwrap());
        }
    }

    #[test] fn pushes_and_pops() {
        let opcodes = vec![Opcode::Pop, Opcode::Push(Json::String("foo".to_string()))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::String("foo".to_string()), vm.run().unwrap());
    }

    #[test] fn dups_tos() {
        let opcodes = vec![Opcode::Dup, Opcode::Pop];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::I64(1), vm.run().unwrap());
    }

    #[test] fn dups_null() {
        let opcodes = vec![Opcode::Pop, Opcode::Dup];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_fields() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"a\": 1}").unwrap());
        assert_eq!(Json::U64(1), vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_missing_field() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"b\": 1}").unwrap());
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_field_from_non_object() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::Index(0), Json::U64(0)),
                         (Opcode::Index(1), Json::U64(1)),
                         (Opcode::Index(2), Json::Null),
                         (Opcode::Index(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn extracts_negative_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::NegativeIndex(0), Json::U64(1)),
                         (Opcode::NegativeIndex(1), Json::U64(0)),
                         (Opcode::NegativeIndex(2), Json::Null),
                         (Opcode::NegativeIndex(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn pushes_type_on_stack() {
        let js = Json::from_str("[0, 1]").unwrap();
        let opcodes = vec![Opcode::Type];
        let tests = vec![(Json::String("a".to_string()), Json::String("string".to_string())),
                         (Json::Array(vec![]), Json::String("array".to_string())),
                         (Json::from_str("{}").unwrap(), Json::String("object".to_string())),
                         (Json::I64(1), Json::String("number".to_string())),
                         (Json::Null, Json::String("null".to_string())),
                         (Json::Boolean(true), Json::String("boolean".to_string()))];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, js);
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn can_jump() {
        let opcodes = vec![Opcode::Jump(2),
                           Opcode::Jump(3),
                           Opcode::Jump(1),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn checks_if_jump_is_valid() {
        let opcodes = vec![Opcode::Jump(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn checks_if_jumpz_is_valid() {
        let opcodes = vec![Opcode::Jumpz(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn can_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Index(0),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(false));
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn can_not_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Jump(100),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(true));
        assert_eq!(Err("Invalid jump address 100".to_string()), vm.run());
    }
}
<commit_msg>Adding Cmp opcode<commit_after>//! Executes compiled JMESPath bytecode instructions.

extern crate rustc_serialize;

use self::rustc_serialize::json::Json;

use ast::{Ast, Comparator, KeyValuePair};

#[derive(Clone, PartialEq, Debug)]
pub enum Opcode {
    // Pops N operands from stack and calls a built-in function by name.
    Call(u8, String),
    // Pops two elements from the stack and pushes true/false.
    Cmp(Comparator),
    // Duplicates the top of the stack
    Dup,
    // Pops TOS and iterates over each element if it is an array. If it
    // is not an array, then jumps to the given address.
    Each(usize),
    // Pops TOS and pushes the <operand> key value of TOS onto the stack.
    // If TOS is not an "object", null is pushed onto the stack.
    Field(String),
    // Pops TOS and pushes the <operand> index value of TOS onto the stack.
    // If TOS is not an "array", null is pushed onto the stack.
    Index(usize),
    // Pops TOS and pushes the <operand> index from the end of the length of
    // of TOS onto the stack. If TOS is not an "array", null is pushed onto
    /// the stack.
    NegativeIndex(usize),
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "array". Pushes T1 onto T2 and pushes T2 back onto the stack.
    InjectIdx(usize),
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "object". Adds T1 at the given operand key onto T2 and pushes T2
    // back onto the stack.
    InjectKey(String),
    // Unconditional jump to an address.
    Jump(usize),
    // Pops the TOS and jumps if 0 or false
    Jumpz(usize),
    // Pops the call stack
    Ret(usize),
    // Pops the TOS
    Pop,
    // Pushes a value onto the stack
    Push(Json),
    // Pops TOS and pushes true/false if the value is truthy.
    Truthy,
    // Pops TOS and pushes the type of the popped value.
    Type
}

pub type Program = Vec<Opcode>;

/// JMESPath VM
pub struct Vm<'a> {
    /// The compiled program.
    opcodes: &'a Program,
    /// The opcode index.
    index: usize,
    /// The VM stack consisting of JSON typed values.
    stack: Vec<Json>
}

impl<'a> Vm<'a> {
    /// Creates a new VM using the given program and data.
    pub fn new(opcodes: &'a Vec<Opcode>, data: Json) -> Vm<'a> {
        Vm {
            opcodes: opcodes,
            stack: vec![data],
            index: 0
        }
    }

    pub fn run(&mut self) -> Result<Json, String> {
        macro_rules! tos {
            () => {{
                self.stack.pop().unwrap_or(Json::Null)
            }};
        }
        while self.index < self.opcodes.len() {
            match &self.opcodes[self.index] {
                &Opcode::Push(ref j) => {
                    self.stack.push(j.clone());
                },
                &Opcode::Pop => {
                    self.stack.pop();
                },
                &Opcode::Dup => {
                    if self.stack.is_empty() {
                        self.stack.push(Json::Null);
                    } else {
                        let tos = self.stack.pop().unwrap();
                        self.stack.push(tos.clone());
                        self.stack.push(tos);
                    }
                },
                &Opcode::Field(ref f) => {
                    match tos!().find(f) {
                        Some(j) => self.stack.push(j.clone()),
                        _ => self.stack.push(Json::Null)
                    };
                },
                &Opcode::Index(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|_| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::NegativeIndex(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|i| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::Truthy => {
                    let tos = tos!();
                    self.stack.push(Json::Boolean(match tos {
                        Json::Boolean(b) if b == false => false,
                        Json::String(ref s) if s.len() == 0 => false,
                        Json::Array(ref a) if a.is_empty() => false,
                        Json::Object(ref o) if o.is_empty() => false,
                        Json::Null => false,
                        _ => true
                    }));
                },
                &Opcode::Type => {
                    let tos = tos!();
                    self.stack.push(match tos {
                        Json::Boolean(_) => Json::String("boolean".to_string()),
                        Json::String(_) => Json::String("string".to_string()),
                        Json::I64(_) | Json::U64(_) | Json::F64(_) => {
                            Json::String("number".to_string())
                        },
                        Json::Array(_) => Json::String("array".to_string()),
                        Json::Object(_) => Json::String("object".to_string()),
                        Json::Null => Json::String("null".to_string()),
                    });
                },
                &Opcode::Jump(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));
                    }
                    self.index = address;
                    continue;
                },
                &Opcode::Jumpz(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));
                    }
                    let tos = tos!();
                    self.index = match tos {
                        Json::Boolean(b) if b == false => address,
                        Json::I64(d) if d == 0 => address,
                        Json::F64(f) if f == 0f64 => address,
                        Json::U64(u) if u == 0 => address,
                        _ => self.index + 1
                    };
                    continue;
                },
                &Opcode::Cmp(ref cmp) => {
                    let (a, b) = (tos!(), tos!());
                    self.stack.push(Json::Boolean(match cmp {
                        &Comparator::Lt
                        | &Comparator::Lte
                        | &Comparator::Gt
                        | &Comparator::Gte => Vm::cmp_numbers(cmp, a, b),
                        &Comparator::Eq => a == b,
                        &Comparator::Ne => a != b,
                    }));
                },
                _ => panic!("Not implemented yet!")
            }
            self.index += 1;
        }

        Ok(self.stack.pop().unwrap_or(Json::Null))
    }

    /// Compares two JSON values, ensuring that each is a number.
    ///
    /// If the provided values are not numbers, then false is returned.
    /// Otherwise, a comparison is made between the two values.
    fn cmp_numbers(cmp: &Comparator, a: Json, b: Json) -> bool {
        if !a.is_number() || !b.is_number() {
            false
        } else {
            let (af, bf) = (a.as_f64().unwrap(), b.as_f64().unwrap());
            match cmp {
                &Comparator::Lt => af < bf,
                &Comparator::Lte => af <= bf,
                &Comparator::Gt => af > bf,
                &Comparator::Gte => af >= bf,
                _ => false
            }
        }
    }
}

#[cfg(test)]
mod test {
    extern crate rustc_serialize;
    use super::{Opcode, Vm};
    use ast::{Comparator};
    use self::rustc_serialize::json::Json;

    #[test] fn truthy_test() {
        let opcodes = vec![Opcode::Truthy];
        let tests = vec![("\"foo\"", true), ("\"\"", false),
                         ("[1]", true), ("[]", false),
                         ("{\"foo\": 1}", true), ("{}", false),
                         ("true", true), ("false", false),
                         ("null", false), ("1", true)];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, Json::from_str(js).unwrap());
            assert_eq!(Json::Boolean(result), vm.run().unwrap());
        }
    }

    #[test] fn pushes_and_pops() {
        let opcodes = vec![Opcode::Pop, Opcode::Push(Json::String("foo".to_string()))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::String("foo".to_string()), vm.run().unwrap());
    }

    #[test] fn dups_tos() {
        let opcodes = vec![Opcode::Dup, Opcode::Pop];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::I64(1), vm.run().unwrap());
    }

    #[test] fn dups_null() {
        let opcodes = vec![Opcode::Pop, Opcode::Dup];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_fields() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"a\": 1}").unwrap());
        assert_eq!(Json::U64(1), vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_missing_field() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"b\": 1}").unwrap());
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_field_from_non_object() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::Index(0), Json::U64(0)),
                         (Opcode::Index(1), Json::U64(1)),
                         (Opcode::Index(2), Json::Null),
                         (Opcode::Index(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn extracts_negative_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::NegativeIndex(0), Json::U64(1)),
                         (Opcode::NegativeIndex(1), Json::U64(0)),
                         (Opcode::NegativeIndex(2), Json::Null),
                         (Opcode::NegativeIndex(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn pushes_type_on_stack() {
        let js = Json::from_str("[0, 1]").unwrap();
        let opcodes = vec![Opcode::Type];
        let tests = vec![(Json::String("a".to_string()), Json::String("string".to_string())),
                         (Json::Array(vec![]), Json::String("array".to_string())),
                         (Json::from_str("{}").unwrap(), Json::String("object".to_string())),
                         (Json::I64(1), Json::String("number".to_string())),
                         (Json::Null, Json::String("null".to_string())),
                         (Json::Boolean(true), Json::String("boolean".to_string()))];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, js);
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn can_jump() {
        let opcodes = vec![Opcode::Jump(2),
                           Opcode::Jump(3),
                           Opcode::Jump(1),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn checks_if_jump_is_valid() {
        let opcodes = vec![Opcode::Jump(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn checks_if_jumpz_is_valid() {
        let opcodes = vec![Opcode::Jumpz(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn can_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Index(0),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(false));
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn can_not_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Jump(100),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(true));
        assert_eq!(Err("Invalid jump address 100".to_string()), vm.run());
    }

    #[test] fn compares_conditionally() {
        let tests = vec![("[0, 1]", Comparator::Lt, false),
                         ("[0, 1]", Comparator::Lte, false),
                         ("[0, 1]", Comparator::Gt, true),
                         ("[0, 1]", Comparator::Gte, true),
                         ("[0, 1]", Comparator::Eq, false),
                         ("[0, \"a\"]", Comparator::Ne, true),
                         ("[0, \"a\"]", Comparator::Eq, false),
                         ("[\"a\", \"a\"]", Comparator::Eq, true),
                         ("[\"a\", \"b\"]", Comparator::Eq, false),
                         ("[\"a\", \"b\"]", Comparator::Ne, true)];
        for (js, cmp, result) in tests {
            let parsed = Json::from_str(js).unwrap();
            let opcodes = vec![Opcode::Push(parsed[0].clone()),
                               Opcode::Push(parsed[1].clone()),
                               Opcode::Cmp(cmp)];
            let mut vm = Vm::new(&opcodes, Json::Null);
            assert_eq!(Json::Boolean(result), vm.run().unwrap());
        }
    }
}
<|endoftext|>


<commit_before>//! A parser for lambda expressions

use term::*;
use self::Token::*;
use self::CToken::*;
use self::Error::*;
use self::Expression::*;
pub use term::Notation::*;

/// A type to represent a parsing error.
#[derive(Debug, PartialEq)]
pub enum Error {
    /// an invalid character was encountered
    InvalidCharacter((usize, char)),
    /// the expression is invalid
    InvalidExpression,
    /// the expression is empty
    EmptyExpression
}

#[derive(Debug, PartialEq)]
enum Token {
    /// λ or \
    Lambda,
    /// left parenthesis
    Lparen,
    /// right parenthesis
    Rparen,
    /// a number
    Number(usize)
}

#[derive(Debug, PartialEq)]
enum CToken {
    /// an abstraction with a bound variable
    CLambda(String),
    /// left parenthesis
    CLparen,
    /// right parenthesis
    CRparen,
    /// a variable with an identifier
    CName(String)
}

fn tokenize_dbr(input: &str) -> Result<Vec<Token>, Error> {
    let chars = input.chars().enumerate();
    let mut tokens = Vec::new();

    for (i, c) in chars {
        match c {
     '\\' | 'λ' => { tokens.push(Lambda) },
            '(' => { tokens.push(Lparen) },
            ')' => { tokens.push(Rparen) },
             _  => {
                if let Some(n) = c.to_digit(16) {
                    tokens.push(Number(n as usize))
                } else if c.is_whitespace() {
                    ()
                } else {
                    return Err(InvalidCharacter((i + 1, c)))
                }
            }
        }
    }

    Ok(tokens)
}

fn tokenize_cla(input: &str) -> Result<Vec<CToken>, Error> {
    let mut chars = input.chars().enumerate().peekable();
    let mut tokens = Vec::new();
    let valid_chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";

    while let Some((i, c)) = chars.next() {
        match c {
            '\\' | 'λ' => {
                let mut name = String::new();
                while let Some((i, c)) = chars.next() {
                    if c == '.' {
                        break
                    } else if valid_chars.contains(c) {
                        name.push(c)
                    } else {
                        return Err(InvalidCharacter((i + 1, c)))
                    }
                }
                tokens.push(CLambda(name))
            },
            '(' => { tokens.push(CLparen) },
            ')' => { tokens.push(CRparen) },
             _  => {
                if c.is_whitespace() {
                    ()
                } else if valid_chars.contains(c) {
                    let mut name = c.to_string();
                    while let Some(&(_, c)) = chars.peek() {
                        if c.is_whitespace() || c == '(' || c == ')' {
                            break
                        } else {
                            name.push(c);
                            chars.next();
                        }
                    }
                    tokens.push(CName(name))
                } else {
                    return Err(InvalidCharacter((i + 1, c)))
                }
            }
        }
    }

    Ok(tokens)
}

fn convert_classic_tokens(tokens: &[CToken]) -> Vec<Token> {
    _convert_classic_tokens(tokens, &mut Vec::new(), &mut 0)
}

fn _convert_classic_tokens(tokens: &[CToken], stack: &mut Vec<String>, pos: &mut usize) -> Vec<Token>
{
    let mut output = Vec::with_capacity(tokens[*pos..].len());
    let mut inner_stack_count = 0;

    while *pos < tokens.len() {
        match tokens[*pos] {
            CLambda(ref name) => {
                output.push(Lambda);
                stack.push(name.clone());
                inner_stack_count += 1;
            },
            CLparen => {
                output.push(Lparen);
                *pos += 1;
                output.append(&mut _convert_classic_tokens(tokens, stack, pos));
            },
            CRparen => {
                output.push(Rparen);
                for _ in 0..inner_stack_count { stack.pop(); }
                return output
            },
            CName(ref name) => {
                if let Some(index) = stack.iter().rev().position(|t| t == name) {
                    output.push(Number(index + 1))
                } else {
                    output.push(Number(stack.len() + 1))
                }
            }
        }
        *pos += 1;
    }

    output
}

#[derive(Debug, PartialEq)]
enum Expression {
    /// an abstraction
    Abstraction,
    /// a sequence of Expressions
    Sequence(Vec<Expression>),
    /// a variable with a De Bruijn index
    Variable(usize)
}

fn get_ast(tokens: &[Token]) -> Result<Expression, Error> {
    _get_ast(tokens, &mut 0)
}

fn _get_ast(tokens: &[Token], pos: &mut usize) -> Result<Expression, Error> {
    let mut expr = Vec::new();

    if tokens.is_empty() { return Err(EmptyExpression) }

    while *pos < tokens.len() {
        match tokens[*pos] {
            Lambda => {
                expr.push(Abstraction)
            },
            Number(i) => {
                expr.push(Variable(i))
            },
            Lparen => {
                *pos += 1;
                let subtree = try!(_get_ast(tokens, pos));
                expr.push(subtree);
            },
            Rparen => {
                return Ok(Sequence(expr))
            }
        }
        *pos += 1;
    }

    Ok(Sequence(expr))
}

/// Parses the input `&str` as a lambda `Term`.
///
/// - lambdas can be represented either with the greek letter (λ) or a backslash (\\ -
/// less aesthetic, but only one byte in size)
/// - the identifiers in `Classic` mode are `String`s of ASCII alphabetic characters
/// - `Classic` mode ignores whitespaces where unambiguous
/// - `DeBruijn` mode ignores all whitespaces (since indices > 15 are very unlikely)
/// - the indices in the `DeBruijn` notation mode start with 1 and are hexadecimal digits
///
/// # Example
/// ```
/// use lambda_calculus::parser::*;
/// use lambda_calculus::church::numerals::{succ, pred};
///
/// assert_eq!(parse(&"λn. λf. λx. n (λg. λh. h (g f)) (λu. x) (λu. u)", Classic), Ok(pred()));
/// assert_eq!(parse(&"λn.λf.λx.n(λg.λh.h(g f))(λu.x)(λu.u)", Classic), Ok(pred()));
///
/// assert_eq!(parse(  &"λλλ2(321)",  DeBruijn), Ok(succ()));
/// assert_eq!(parse(&r#"\\\2(321)"#, DeBruijn), Ok(succ()));
/// ```
pub fn parse(input: &str, notation: Notation) -> Result<Term, Error> {
    let tokens = if notation == DeBruijn {
        try!(tokenize_dbr(input))
    } else {
        convert_classic_tokens(&try!(tokenize_cla(input)))
    };
    let ast = try!(get_ast(&tokens));

    let exprs = try!(if let Sequence(exprs) = ast { Ok(exprs) } else { Err(InvalidExpression) });

    fold_exprs(&exprs, &mut Vec::new(), &mut Vec::new())
}

fn fold_exprs(exprs: &[Expression], stack: &mut Vec<Expression>, output: &mut Vec<Term>)
    -> Result<Term, Error>
{
    let iter = exprs.into_iter();

    for expr in iter {
        match *expr {
            Variable(i) => output.push(Var(i)),
            Abstraction => stack.push(Abstraction),
            Sequence(ref exprs) => {
                let subexpr = try!(fold_exprs(exprs, &mut Vec::new(), &mut Vec::new()));
                output.push(subexpr)
            }
        }
    }

    let mut ret = try!(fold_terms(output.drain(..).collect()));

    while let Some(Abstraction) = stack.pop() {
        ret = abs(ret);
    }

    Ok(ret)
}

fn fold_terms(mut terms: Vec<Term>) -> Result<Term, Error> {
    if terms.len() > 1 {
        let fst = terms.remove(0);
        Ok( terms.into_iter().fold(fst, app) )
    } else if terms.len() == 1 {
        Ok( terms.pop().unwrap() ) // safe; ensured above
    } else {
        Err(EmptyExpression)
    }
}

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

    #[test]
    fn tokenization_error() {
        assert_eq!(tokenize_dbr(&"λλx2"),    Err(InvalidCharacter((3, 'x'))));
        assert_eq!(tokenize_cla(&"λa.λb a"), Err(InvalidCharacter((6, ' '))));
    }

    #[test]
    fn tokenization_success() {
        let quine = "λ 1 ( (λ 1 1) (λ λ λ λ λ 1 4 (3 (5 5) 2) ) ) 1";
        let tokens = tokenize_dbr(&quine);

        assert!(tokens.is_ok());
        assert_eq!(tokens.unwrap(), vec![Lambda, Number(1), Lparen, Lparen, Lambda, Number(1),
            Number(1), Rparen, Lparen, Lambda, Lambda, Lambda, Lambda, Lambda, Number(1),
            Number(4), Lparen, Number(3), Lparen, Number(5), Number(5), Rparen, Number(2),
            Rparen, Rparen, Rparen, Number(1)]);
    }

    #[test]
    fn tokenization_success_classic() {
        let blc_dbr = "(λ11)(λλλ1(λλλλ3(λ5(3(λ2(3(λλ3(λ123)))(4(λ4(λ31(21))))))(1(2(λ12))\
            (λ4(λ4(λ2(14)))5))))(33)2)(λ1((λ11)(λ11)))";
        let blc_cla = format!("{}", parse(&blc_dbr, DeBruijn).unwrap());

        let tokens_cla = tokenize_cla(&blc_cla);
        let tokens_dbr = tokenize_dbr(&blc_dbr);

        assert!(tokens_cla.is_ok());
        assert!(tokens_dbr.is_ok());

        assert_eq!(convert_classic_tokens(&tokens_cla.unwrap()), tokens_dbr.unwrap());
    }

    #[test]
    fn alternative_lambda_parsing() {
        assert_eq!(parse(&r#"\\\2(321)"#, DeBruijn), parse(&"λλλ2(321)", DeBruijn))
    }

    #[test]
    fn succ_ast() {
        let tokens = tokenize_dbr(&"λλλ2(321)").unwrap();
        let ast = get_ast(&tokens);

        assert_eq!(ast,
            Ok(Sequence(vec![
                Abstraction,
                Abstraction,
                Abstraction,
                Variable(2),
                Sequence(vec![
                    Variable(3),
                    Variable(2),
                    Variable(1)
                ])
            ])
        ));
    }

    #[test]
    fn parse_y() {
        let y = "λ(λ2(11))(λ2(11))";
        assert_eq!(parse(&y, DeBruijn).unwrap(),
            abs(
                app(
                    abs(app(Var(2), app(Var(1), Var(1)))),
                    abs(app(Var(2), app(Var(1), Var(1))))
                )
            )
        );
    }

    #[test]
    fn parse_quine() {
        let quine = "λ1((λ11)(λλλλλ14(3(55)2)))1";
        assert_eq!(parse(&quine, DeBruijn).unwrap(),
            abs(app(app(Var(1), app(abs(app(Var(1), Var(1))), abs(abs(abs(abs(abs(app(app(Var(1),
            Var(4)), app(app(Var(3), app(Var(5), Var(5))), Var(2)))))))))), Var(1)))
        );
    }

    #[test]
    fn parse_blc() {
        let blc = "(λ11)(λλλ1(λλλλ3(λ5(3(λ2(3(λλ3(λ123)))(4(λ4(λ31(21))))))(1(2(λ12))\
                   (λ4(λ4(λ2(14)))5))))(33)2)(λ1((λ11)(λ11)))";
        assert_eq!(parse(&blc, DeBruijn).unwrap(),
            app(app(abs(app(Var(1), Var(1))), abs(abs(abs(app(app(app(Var(1),
            abs(abs(abs(abs(app(Var(3), abs(app(app(Var(5), app(Var(3), abs(app(app(Var(2),
            app(Var(3), abs(abs(app(Var(3), abs(app(app(Var(1), Var(2)), Var(3)))))))), app(Var(4),
            abs(app(Var(4), abs(app(app(Var(3), Var(1)), app(Var(2), Var(1))))))))))),
            app(app(Var(1), app(Var(2), abs(app(Var(1), Var(2))))), abs(app(app(Var(4),
            abs(app(Var(4), abs(app(Var(2), app(Var(1), Var(4))))))), Var(5)))))))))))),
            app(Var(3), Var(3))), Var(2)))))), abs(app(Var(1), app(abs(app(Var(1), Var(1))),
            abs(app(Var(1), Var(1)))))))
        );
    }
}
<commit_msg>improve fold_exprs() performance<commit_after>//! A parser for lambda expressions

use term::*;
use self::Token::*;
use self::CToken::*;
use self::Error::*;
use self::Expression::*;
pub use term::Notation::*;

/// A type to represent a parsing error.
#[derive(Debug, PartialEq)]
pub enum Error {
    /// an invalid character was encountered
    InvalidCharacter((usize, char)),
    /// the expression is invalid
    InvalidExpression,
    /// the expression is empty
    EmptyExpression
}

#[derive(Debug, PartialEq)]
enum Token {
    /// λ or \
    Lambda,
    /// left parenthesis
    Lparen,
    /// right parenthesis
    Rparen,
    /// a number
    Number(usize)
}

#[derive(Debug, PartialEq)]
enum CToken {
    /// an abstraction with a bound variable
    CLambda(String),
    /// left parenthesis
    CLparen,
    /// right parenthesis
    CRparen,
    /// a variable with an identifier
    CName(String)
}

fn tokenize_dbr(input: &str) -> Result<Vec<Token>, Error> {
    let chars = input.chars().enumerate();
    let mut tokens = Vec::new();

    for (i, c) in chars {
        match c {
     '\\' | 'λ' => { tokens.push(Lambda) },
            '(' => { tokens.push(Lparen) },
            ')' => { tokens.push(Rparen) },
             _  => {
                if let Some(n) = c.to_digit(16) {
                    tokens.push(Number(n as usize))
                } else if c.is_whitespace() {
                    ()
                } else {
                    return Err(InvalidCharacter((i + 1, c)))
                }
            }
        }
    }

    Ok(tokens)
}

fn tokenize_cla(input: &str) -> Result<Vec<CToken>, Error> {
    let mut chars = input.chars().enumerate().peekable();
    let mut tokens = Vec::new();
    let valid_chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";

    while let Some((i, c)) = chars.next() {
        match c {
            '\\' | 'λ' => {
                let mut name = String::new();
                while let Some((i, c)) = chars.next() {
                    if c == '.' {
                        break
                    } else if valid_chars.contains(c) {
                        name.push(c)
                    } else {
                        return Err(InvalidCharacter((i + 1, c)))
                    }
                }
                tokens.push(CLambda(name))
            },
            '(' => { tokens.push(CLparen) },
            ')' => { tokens.push(CRparen) },
             _  => {
                if c.is_whitespace() {
                    ()
                } else if valid_chars.contains(c) {
                    let mut name = c.to_string();
                    while let Some(&(_, c)) = chars.peek() {
                        if c.is_whitespace() || c == '(' || c == ')' {
                            break
                        } else {
                            name.push(c);
                            chars.next();
                        }
                    }
                    tokens.push(CName(name))
                } else {
                    return Err(InvalidCharacter((i + 1, c)))
                }
            }
        }
    }

    Ok(tokens)
}

fn convert_classic_tokens(tokens: &[CToken]) -> Vec<Token> {
    _convert_classic_tokens(tokens, &mut Vec::new(), &mut 0)
}

fn _convert_classic_tokens(tokens: &[CToken], stack: &mut Vec<String>, pos: &mut usize) -> Vec<Token>
{
    let mut output = Vec::with_capacity(tokens[*pos..].len());
    let mut inner_stack_count = 0;

    while *pos < tokens.len() {
        match tokens[*pos] {
            CLambda(ref name) => {
                output.push(Lambda);
                stack.push(name.clone());
                inner_stack_count += 1;
            },
            CLparen => {
                output.push(Lparen);
                *pos += 1;
                output.append(&mut _convert_classic_tokens(tokens, stack, pos));
            },
            CRparen => {
                output.push(Rparen);
                for _ in 0..inner_stack_count { stack.pop(); }
                return output
            },
            CName(ref name) => {
                if let Some(index) = stack.iter().rev().position(|t| t == name) {
                    output.push(Number(index + 1))
                } else {
                    output.push(Number(stack.len() + 1))
                }
            }
        }
        *pos += 1;
    }

    output
}

#[derive(Debug, PartialEq)]
enum Expression {
    /// an abstraction
    Abstraction,
    /// a sequence of Expressions
    Sequence(Vec<Expression>),
    /// a variable with a De Bruijn index
    Variable(usize)
}

fn get_ast(tokens: &[Token]) -> Result<Expression, Error> {
    _get_ast(tokens, &mut 0)
}

fn _get_ast(tokens: &[Token], pos: &mut usize) -> Result<Expression, Error> {
    let mut expr = Vec::new();

    if tokens.is_empty() { return Err(EmptyExpression) }

    while *pos < tokens.len() {
        match tokens[*pos] {
            Lambda => {
                expr.push(Abstraction)
            },
            Number(i) => {
                expr.push(Variable(i))
            },
            Lparen => {
                *pos += 1;
                let subtree = try!(_get_ast(tokens, pos));
                expr.push(subtree);
            },
            Rparen => {
                return Ok(Sequence(expr))
            }
        }
        *pos += 1;
    }

    Ok(Sequence(expr))
}

/// Parses the input `&str` as a lambda `Term`.
///
/// - lambdas can be represented either with the greek letter (λ) or a backslash (\\ -
/// less aesthetic, but only one byte in size)
/// - the identifiers in `Classic` mode are `String`s of ASCII alphabetic characters
/// - `Classic` mode ignores whitespaces where unambiguous
/// - `DeBruijn` mode ignores all whitespaces (since indices > 15 are very unlikely)
/// - the indices in the `DeBruijn` notation mode start with 1 and are hexadecimal digits
///
/// # Example
/// ```
/// use lambda_calculus::parser::*;
/// use lambda_calculus::church::numerals::{succ, pred};
///
/// assert_eq!(parse(&"λn. λf. λx. n (λg. λh. h (g f)) (λu. x) (λu. u)", Classic), Ok(pred()));
/// assert_eq!(parse(&"λn.λf.λx.n(λg.λh.h(g f))(λu.x)(λu.u)", Classic), Ok(pred()));
///
/// assert_eq!(parse(  &"λλλ2(321)",  DeBruijn), Ok(succ()));
/// assert_eq!(parse(&r#"\\\2(321)"#, DeBruijn), Ok(succ()));
/// ```
pub fn parse(input: &str, notation: Notation) -> Result<Term, Error> {
    let tokens = if notation == DeBruijn {
        try!(tokenize_dbr(input))
    } else {
        convert_classic_tokens(&try!(tokenize_cla(input)))
    };
    let ast = try!(get_ast(&tokens));

    let exprs = try!(if let Sequence(exprs) = ast { Ok(exprs) } else { Err(InvalidExpression) });

    fold_exprs(&exprs)
}

fn fold_exprs(exprs: &[Expression]) -> Result<Term, Error> {
    let iter = exprs.into_iter();
    let mut stack  = Vec::new();
    let mut output = Vec::new();

    for expr in iter {
        match *expr {
            Variable(i) => output.push(Var(i)),
            Abstraction => stack.push(Abstraction),
            Sequence(ref exprs) => {
                let subexpr = try!(fold_exprs(exprs));
                output.push(subexpr)
            }
        }
    }

    let mut ret = try!(fold_terms(output));

    while let Some(Abstraction) = stack.pop() {
        ret = abs(ret);
    }

    Ok(ret)
}

fn fold_terms(mut terms: Vec<Term>) -> Result<Term, Error> {
    if terms.len() > 1 {
        let fst = terms.remove(0);
        Ok( terms.into_iter().fold(fst, app) )
    } else if terms.len() == 1 {
        Ok( terms.pop().unwrap() ) // safe; ensured above
    } else {
        Err(EmptyExpression)
    }
}

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

    #[test]
    fn tokenization_error() {
        assert_eq!(tokenize_dbr(&"λλx2"),    Err(InvalidCharacter((3, 'x'))));
        assert_eq!(tokenize_cla(&"λa.λb a"), Err(InvalidCharacter((6, ' '))));
    }

    #[test]
    fn tokenization_success() {
        let quine = "λ 1 ( (λ 1 1) (λ λ λ λ λ 1 4 (3 (5 5) 2) ) ) 1";
        let tokens = tokenize_dbr(&quine);

        assert!(tokens.is_ok());
        assert_eq!(tokens.unwrap(), vec![Lambda, Number(1), Lparen, Lparen, Lambda, Number(1),
            Number(1), Rparen, Lparen, Lambda, Lambda, Lambda, Lambda, Lambda, Number(1),
            Number(4), Lparen, Number(3), Lparen, Number(5), Number(5), Rparen, Number(2),
            Rparen, Rparen, Rparen, Number(1)]);
    }

    #[test]
    fn tokenization_success_classic() {
        let blc_dbr = "(λ11)(λλλ1(λλλλ3(λ5(3(λ2(3(λλ3(λ123)))(4(λ4(λ31(21))))))(1(2(λ12))\
            (λ4(λ4(λ2(14)))5))))(33)2)(λ1((λ11)(λ11)))";
        let blc_cla = format!("{}", parse(&blc_dbr, DeBruijn).unwrap());

        let tokens_cla = tokenize_cla(&blc_cla);
        let tokens_dbr = tokenize_dbr(&blc_dbr);

        assert!(tokens_cla.is_ok());
        assert!(tokens_dbr.is_ok());

        assert_eq!(convert_classic_tokens(&tokens_cla.unwrap()), tokens_dbr.unwrap());
    }

    #[test]
    fn alternative_lambda_parsing() {
        assert_eq!(parse(&r#"\\\2(321)"#, DeBruijn), parse(&"λλλ2(321)", DeBruijn))
    }

    #[test]
    fn succ_ast() {
        let tokens = tokenize_dbr(&"λλλ2(321)").unwrap();
        let ast = get_ast(&tokens);

        assert_eq!(ast,
            Ok(Sequence(vec![
                Abstraction,
                Abstraction,
                Abstraction,
                Variable(2),
                Sequence(vec![
                    Variable(3),
                    Variable(2),
                    Variable(1)
                ])
            ])
        ));
    }

    #[test]
    fn parse_y() {
        let y = "λ(λ2(11))(λ2(11))";
        assert_eq!(parse(&y, DeBruijn).unwrap(),
            abs(
                app(
                    abs(app(Var(2), app(Var(1), Var(1)))),
                    abs(app(Var(2), app(Var(1), Var(1))))
                )
            )
        );
    }

    #[test]
    fn parse_quine() {
        let quine = "λ1((λ11)(λλλλλ14(3(55)2)))1";
        assert_eq!(parse(&quine, DeBruijn).unwrap(),
            abs(app(app(Var(1), app(abs(app(Var(1), Var(1))), abs(abs(abs(abs(abs(app(app(Var(1),
            Var(4)), app(app(Var(3), app(Var(5), Var(5))), Var(2)))))))))), Var(1)))
        );
    }

    #[test]
    fn parse_blc() {
        let blc = "(λ11)(λλλ1(λλλλ3(λ5(3(λ2(3(λλ3(λ123)))(4(λ4(λ31(21))))))(1(2(λ12))\
                   (λ4(λ4(λ2(14)))5))))(33)2)(λ1((λ11)(λ11)))";
        assert_eq!(parse(&blc, DeBruijn).unwrap(),
            app(app(abs(app(Var(1), Var(1))), abs(abs(abs(app(app(app(Var(1),
            abs(abs(abs(abs(app(Var(3), abs(app(app(Var(5), app(Var(3), abs(app(app(Var(2),
            app(Var(3), abs(abs(app(Var(3), abs(app(app(Var(1), Var(2)), Var(3)))))))), app(Var(4),
            abs(app(Var(4), abs(app(app(Var(3), Var(1)), app(Var(2), Var(1))))))))))),
            app(app(Var(1), app(Var(2), abs(app(Var(1), Var(2))))), abs(app(app(Var(4),
            abs(app(Var(4), abs(app(Var(2), app(Var(1), Var(4))))))), Var(5)))))))))))),
            app(Var(3), Var(3))), Var(2)))))), abs(app(Var(1), app(abs(app(Var(1), Var(1))),
            abs(app(Var(1), Var(1)))))))
        );
    }
}
<|endoftext|>


<commit_before>use term::*;
use term::Term::*;

fn _parse(input: &str, depth: usize) -> Option<(Term, &str)> {
    let c = input.chars().next();

    match c {
        Some('λ') => {
            if let Some((term, rest)) = _parse(&input[2..], depth) {
                Some((abs(term), rest))
            } else {
                None
            }
        },
        Some('(') => {
            let range = input[1..].chars().take_while(|&c| c != ')');
            Some((app(parse(&input[1..range]), parse(&input[range..])), &input))
        },
        Some(')') => _parse(&input[1..], depth),
        Some(x) => Some((Var(x.to_digit(16).unwrap() as usize), &input[1..])),
        _ => None
    }
}

pub fn parse(input: &str) -> Term {
    _parse(input, 0).unwrap().0
}

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

    #[test]
    fn parsing() {
        let tru = "λλ2";
        let quine = "λ1((λ11)(λλλλλ14(3(55)2)))1";

        assert_eq!(parse(tru), abs(abs(Var(2))));
        assert_eq!(format!("{}", parse(quine)), quine);
    }
}
<commit_msg>Shunting-yardy parsing approach<commit_after>#![allow(dead_code)]

use term::*;
use term::Term::*;
use self::Token::*;
use self::Error::*;

#[derive(Debug, PartialEq)]
enum Error {
    InvalidCharacter((usize, char))
}

#[derive(Debug, PartialEq, Clone, Copy)]
enum Token {
    Lambda,
    Lparen,
    Rparen,
    Index(usize)
}

#[derive(Debug)]
struct Lexer {
    tokens: Vec<Token>,
    position: usize
}

impl Lexer {
    fn new() -> Lexer {
        Lexer { tokens: Vec::new(), position: 0 }
    }

    fn tokenize(&mut self, input: &str) -> Result<(), Error> {
        let mut chars = input.chars();

        while let Some(c) = chars.next() {
            match c {
                'λ' => { self.tokens.push(Lambda) },
                '(' => { self.tokens.push(Lparen) },
                ')' => { self.tokens.push(Rparen) },
                 x  => {
                    if let Some(i) = x.to_digit(16) {
                        self.tokens.push(Index(i as usize))
                    } else if x.is_whitespace() {
                        ()
                    } else {
                        return Err(InvalidCharacter((self.position, x)))
                    }
                }
            }
            self.position += if c == 'λ' { 2 } else { 1 };
        }

        Ok(())
    }
}

#[derive(Debug)]
struct Parser {
    output: Vec<Term>,
    stack: Vec<Token>,
    temp: Vec<Term>
}

impl Parser {
    fn new() -> Parser {
        Parser { output: Vec::new(), stack: Vec::new(), temp: Vec::new() }
    }

    fn process(&mut self, tokens: &[Token]) -> Result<Term, Error> {
        let mut iter = tokens.iter();

        while let Some(&token) = iter.next() {
            println!("stack: {:?}; output: {:?}; temp: {:?}", self.stack, self.output, self.temp);
            match token {
                Lambda => {
                    self.stack.push(token)
                },
                Lparen => {
                    while !self.temp.is_empty() {
                        self.output.push(self.temp.pop().unwrap())
                    }
                    self.stack.push(token)
                },
                Rparen => {
                    if !self.temp.is_empty() {

                        let mut ret = fold_terms(self.temp.clone());
                        while self.stack.last() == Some(&Lambda) {
                            ret = abs(ret);
                        }
                        self.output.push(ret);

//                        self.output.push(fold_terms(self.temp.clone())); // TODO: some swap or other smartery
                        self.temp.clear();
                    }
                    self.stack.pop(); // drop Lparen; TODO: check if Lparen
                },
                Index(i) => {
                    if self.stack.contains(&Lparen) {
                        self.temp.push(Var(i))
                    } else {
                        self.output.push(Var(i))
                    }
                }
            }
        }

        let mut ret = fold_terms(self.output.clone());

        while let Some(Lambda) = self.stack.pop() {
            ret = abs(ret);
        }

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

        Ok(ret)
    }

    fn parse(&mut self, string: &str) -> Result<Term, Error> {
        let mut lexer = Lexer::new();
        try!(lexer.tokenize(string));
        let mut parser = Parser::new();
        parser.process(&lexer.tokens)
    }
}

fn fold_terms(mut terms: Vec<Term>) -> Term {
//    println!("folding terms {:?}", terms);
    if terms.len() > 1 {
        terms.reverse();
        let fst = terms.pop().unwrap();
        terms.reverse();
        terms.into_iter().fold(fst, |acc, t| app(acc, t))
    } else {
        terms.pop().unwrap()
    }
}
/*
pub fn split_parenthesized(input: &str) -> (&str, &str) {
    let mut pos = 1;
    let mut depth = 1;
    let mut chars = input.chars().skip(1);

    while let Some(c) = chars.next() {
        if depth == 0 {
            break;
        } else if c == '(' {
            depth += 1;
        } else if c == ')' {
            depth -= 1;
        }
        if c == 'λ' { pos += 2 } else { pos += 1 }
    }
//    println!("input has a parens \"{}\" and rest \"{}\"", &input[1..pos-1], &input[pos..]);
    (&input[1..pos-1], &input[pos..])
}
*/
#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn tokenization_error() {
        let fail = "λλx2";
        let mut lexer = Lexer::new();

        assert_eq!(lexer.tokenize(&fail), Err(InvalidCharacter((4, 'x'))))
    }

    #[test]
    fn tokenization_success() {
        let quine = "λ 1 ( (λ 1 1) (λ λ λ λ λ 1 4 (3 (5 5) 2) ) ) 1";
        let mut lexer = Lexer::new();

        assert!(lexer.tokenize(&quine).is_ok());
        assert_eq!(lexer.tokens, vec![Lambda, Index(1), Lparen, Lparen, Lambda, Index(1), Index(1), Rparen, Lparen, Lambda,
                                      Lambda, Lambda, Lambda, Lambda, Index(1), Index(4), Lparen, Index(3), Lparen, Index(5),
                                      Index(5), Rparen, Index(2), Rparen, Rparen, Rparen, Index(1)]
        )
    }
/*
    #[test]
    fn parse_succ() {
        let succ = "λλλ2(321)";
        let mut parser = Parser::new();

        assert_eq!(&*format!("{}", parser.parse(&succ).unwrap()), "λλλ2(321)");
    }
*/
    #[test]
    fn parse_y() {
        let y = "λ(λ2(11))(λ2(11))";
        let mut parser = Parser::new();

        assert_eq!(&*format!("{}", parser.parse(&y).unwrap()), "λ(λ2(11))(λ2(11))");
    }
}
<|endoftext|>