<commit_before>extern crate bytes;
extern crate env_logger;
extern crate mio;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;
extern crate failure;
extern crate slab;

use clap::{App, Arg};
use failure::Error;
use failure::ResultExt;
use mio::net::{TcpListener, TcpStream};
use slab::Slab;
use std::collections::VecDeque;
use std::io::{Read, Write};
use std::ops;

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

trait RuleHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready);
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token);
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error>;
    fn is_closed(&self) -> bool;
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::Poll, listener: TcpListener) {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Listener(Listener::new(listener, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token);
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::Poll) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => self.process_broadcast(s),
            MiChatCommand::NewConnection(socket) => self.process_new_connection(event_loop, socket),
        }
    }

    fn process_broadcast(&mut self, s: String) {
        for c in self.client_connections_mut() {
            c.enqueue(&s);
        }
    }

    fn process_new_connection(&mut self, event_loop: &mut mio::Poll, socket: mio::tcp::TcpStream) {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Conn(Connection::new(socket, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token);
    }

    fn client_connections_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut Connection> + 'a {
        self.connections.iter_mut().filter_map(|(_, it)| match it {
            &mut EventHandler::Conn(ref mut c) => Some(c),
            _ => None,
        })
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::Ready,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl RuleHandler for Connection {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) {
        event_loop
            .register(
                &self.socket,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).expect("event loop register");
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        self.sock_status.insert(events);
        info!(
            "Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.socket.peer_addr(),
            events,
            self.sock_status
        );
    }

    // actions are processed here on down.

    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            self.sock_status.remove(mio::Ready::readable());
            self.read()?;
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.sock_status.remove(mio::Ready::writable());
            self.write()?;
        }

        if !self.is_closed() {
            self.reregister(event_loop)?;
        }

        Ok(())
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::Ready::empty(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }
    fn process_buffer(&mut self, to_parent: &mut FnMut(MiChatCommand)) {
        let mut prev = 0;
        info!(
            "{:?}: Read buffer: {:?}",
            self.socket.peer_addr(),
            self.read_buf
        );
        for n in self.read_buf.iter().enumerate().filter_map(|(i, e)| {
            if *e == '\n' as u8 {
                Some(i)
            } else {
                None
            }
        }) {
            info!(
                "{:?}: Pos: {:?}; chunk: {:?}",
                self.socket.peer_addr(),
                n,
                &self.read_buf[prev..n]
            );
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent(cmd);
            prev = n + 1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!(
            "{:?}: read Remainder: {}",
            self.socket.peer_addr(),
            remainder.len()
        );
        self.read_buf = remainder;
    }

    fn read(&mut self) -> Result<(), Error> {
        let mut abuf = vec![0; 1024];
        match self.socket.read(&mut abuf).context("Reading socket")? {
            0 => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true;
                Ok(())
            }
            n => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(&abuf[0..n]);
                Ok(())
            }
        }
    }

    fn write(&mut self) -> Result<(), Error> {
        let n = self
            .socket
            .write(&mut self.write_buf)
            .context("writing buffer")?;
        info!(
            "{:?}: Wrote {} of {} in buffer",
            self.socket.peer_addr(),
            n,
            self.write_buf.len()
        );
        self.write_buf = self.write_buf[n..].to_vec();
        info!(
            "{:?}: Now {:?}b",
            self.socket.peer_addr(),
            self.write_buf.len()
        );
        Ok(())
    }

    fn reregister(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut flags = mio::Ready::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::Ready::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop
            .reregister(&self.socket, self.token, flags, mio::PollOpt::oneshot())
            .context("EventLoop#reregister")?;

        Ok(())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::Ready,
    token: mio::Token,
}

impl Listener {
    fn new(listener: TcpListener, token: mio::Token) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::Ready::empty(),
            token: token,
        }
    }
    fn reregister(&self, event_loop: &mut mio::Poll, token: mio::Token) {
        event_loop
            .reregister(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).expect("Register listener");
    }
}

impl RuleHandler for Listener {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) {
        event_loop
            .register(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).expect("Register listener");
    }

    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!(
            "Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.listener.local_addr(),
            events,
            self.sock_status
        );
    }
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            info!("the listener socket is ready to accept a connection");
            self.sock_status.remove(mio::Ready::readable());
            match self.listener.accept().context("Accept listening socket")? {
                (socket, _client) => {
                    let cmd = MiChatCommand::NewConnection(socket);
                    to_parent(cmd);
                }
            }

            self.reregister(event_loop, self.token);
        }
        Ok(())
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener(Listener),
    Conn(Connection),
}

impl ops::Deref for EventHandler {
    type Target = RuleHandler;
    fn deref(&self) -> &Self::Target {
        match self {
            &EventHandler::Conn(ref conn) => conn,
            &EventHandler::Listener(ref listener) => listener,
        }
    }
}
impl ops::DerefMut for EventHandler {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn,
            &mut EventHandler::Listener(ref mut listener) => listener,
        }
    }
}

impl MiChat {
    fn ready(&mut self, event_loop: &mut mio::Poll, token: mio::Token, events: mio::Ready) {
        info!("{:?}: {:?}", token, events);
        self.connections[token.into()].handle_event(event_loop, events);
        if self.connections[token.into()].is_closed() {
            info!("Removing; {:?}", token);
            self.connections.remove(token.into());
        }

        let mut parent_actions = VecDeque::new();
        let mut failed = Vec::new();
        loop {
            for (idx, conn) in self.connections.iter_mut() {
                if let Err(e) =
                    conn.process_rules(event_loop, &mut |action| parent_actions.push_back(action))
                {
                    failed.push(idx);
                    error!("Error on connection {:?}; Dropping: {:?}", idx, e);
                }
            }
            for idx in failed.drain(..) {
                self.connections.remove(idx);
            }
            // Anything left to process?
            if parent_actions.is_empty() {
                break;
            }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }

    fn run(&mut self, event_loop: &mut mio::Poll) {
        loop {
            self.run_once(event_loop);
        }
    }
    fn run_once(&mut self, event_loop: &mut mio::Poll) {
        let mut events = mio::Events::with_capacity(1024);
        event_loop.poll(&mut events, None).expect("Poll");
        for ev in &events {
            self.ready(event_loop, ev.token(), ev.readiness())
        }
    }
}

fn main() {
    env_logger::init();
    let matches = App::new("michat")
        .arg(
            Arg::with_name("bind")
                .short("l")
                .takes_value(true)
                .required(true),
        ).get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::Poll::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    service.run(&mut event_loop);
}
<commit_msg>Errors on reregister listener.<commit_after>extern crate bytes;
extern crate env_logger;
extern crate mio;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;
extern crate failure;
extern crate slab;

use clap::{App, Arg};
use failure::Error;
use failure::ResultExt;
use mio::net::{TcpListener, TcpStream};
use slab::Slab;
use std::collections::VecDeque;
use std::io::{Read, Write};
use std::ops;

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

trait RuleHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready);
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token);
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error>;
    fn is_closed(&self) -> bool;
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::Poll, listener: TcpListener) {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Listener(Listener::new(listener, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token);
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::Poll) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => self.process_broadcast(s),
            MiChatCommand::NewConnection(socket) => self.process_new_connection(event_loop, socket),
        }
    }

    fn process_broadcast(&mut self, s: String) {
        for c in self.client_connections_mut() {
            c.enqueue(&s);
        }
    }

    fn process_new_connection(&mut self, event_loop: &mut mio::Poll, socket: mio::tcp::TcpStream) {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Conn(Connection::new(socket, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token);
    }

    fn client_connections_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut Connection> + 'a {
        self.connections.iter_mut().filter_map(|(_, it)| match it {
            &mut EventHandler::Conn(ref mut c) => Some(c),
            _ => None,
        })
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::Ready,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl RuleHandler for Connection {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) {
        event_loop
            .register(
                &self.socket,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).expect("event loop register");
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        self.sock_status.insert(events);
        info!(
            "Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.socket.peer_addr(),
            events,
            self.sock_status
        );
    }

    // actions are processed here on down.

    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            self.sock_status.remove(mio::Ready::readable());
            self.read()?;
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.sock_status.remove(mio::Ready::writable());
            self.write()?;
        }

        if !self.is_closed() {
            self.reregister(event_loop)?;
        }

        Ok(())
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::Ready::empty(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }
    fn process_buffer(&mut self, to_parent: &mut FnMut(MiChatCommand)) {
        let mut prev = 0;
        info!(
            "{:?}: Read buffer: {:?}",
            self.socket.peer_addr(),
            self.read_buf
        );
        for n in self.read_buf.iter().enumerate().filter_map(|(i, e)| {
            if *e == '\n' as u8 {
                Some(i)
            } else {
                None
            }
        }) {
            info!(
                "{:?}: Pos: {:?}; chunk: {:?}",
                self.socket.peer_addr(),
                n,
                &self.read_buf[prev..n]
            );
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent(cmd);
            prev = n + 1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!(
            "{:?}: read Remainder: {}",
            self.socket.peer_addr(),
            remainder.len()
        );
        self.read_buf = remainder;
    }

    fn read(&mut self) -> Result<(), Error> {
        let mut abuf = vec![0; 1024];
        match self.socket.read(&mut abuf).context("Reading socket")? {
            0 => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true;
                Ok(())
            }
            n => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(&abuf[0..n]);
                Ok(())
            }
        }
    }

    fn write(&mut self) -> Result<(), Error> {
        let n = self
            .socket
            .write(&mut self.write_buf)
            .context("writing buffer")?;
        info!(
            "{:?}: Wrote {} of {} in buffer",
            self.socket.peer_addr(),
            n,
            self.write_buf.len()
        );
        self.write_buf = self.write_buf[n..].to_vec();
        info!(
            "{:?}: Now {:?}b",
            self.socket.peer_addr(),
            self.write_buf.len()
        );
        Ok(())
    }

    fn reregister(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut flags = mio::Ready::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::Ready::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop
            .reregister(&self.socket, self.token, flags, mio::PollOpt::oneshot())
            .context("EventLoop#reregister")?;

        Ok(())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::Ready,
    token: mio::Token,
}

impl Listener {
    fn new(listener: TcpListener, token: mio::Token) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::Ready::empty(),
            token: token,
        }
    }
    fn reregister(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error> {
        event_loop
            .reregister(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).context("Register listener")?;
            Ok(())
    }
}

impl RuleHandler for Listener {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) {
        event_loop
            .register(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).expect("Register listener");
    }

    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!(
            "Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.listener.local_addr(),
            events,
            self.sock_status
        );
    }
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            info!("the listener socket is ready to accept a connection");
            self.sock_status.remove(mio::Ready::readable());
            match self.listener.accept().context("Accept listening socket")? {
                (socket, _client) => {
                    let cmd = MiChatCommand::NewConnection(socket);
                    to_parent(cmd);
                }
            }

            self.reregister(event_loop, self.token)?;
        }
        Ok(())
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener(Listener),
    Conn(Connection),
}

impl ops::Deref for EventHandler {
    type Target = RuleHandler;
    fn deref(&self) -> &Self::Target {
        match self {
            &EventHandler::Conn(ref conn) => conn,
            &EventHandler::Listener(ref listener) => listener,
        }
    }
}
impl ops::DerefMut for EventHandler {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn,
            &mut EventHandler::Listener(ref mut listener) => listener,
        }
    }
}

impl MiChat {
    fn ready(&mut self, event_loop: &mut mio::Poll, token: mio::Token, events: mio::Ready) {
        info!("{:?}: {:?}", token, events);
        self.connections[token.into()].handle_event(event_loop, events);
        if self.connections[token.into()].is_closed() {
            info!("Removing; {:?}", token);
            self.connections.remove(token.into());
        }

        let mut parent_actions = VecDeque::new();
        let mut failed = Vec::new();
        loop {
            for (idx, conn) in self.connections.iter_mut() {
                if let Err(e) =
                    conn.process_rules(event_loop, &mut |action| parent_actions.push_back(action))
                {
                    failed.push(idx);
                    error!("Error on connection {:?}; Dropping: {:?}", idx, e);
                }
            }
            for idx in failed.drain(..) {
                self.connections.remove(idx);
            }
            // Anything left to process?
            if parent_actions.is_empty() {
                break;
            }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }

    fn run(&mut self, event_loop: &mut mio::Poll) {
        loop {
            self.run_once(event_loop);
        }
    }
    fn run_once(&mut self, event_loop: &mut mio::Poll) {
        let mut events = mio::Events::with_capacity(1024);
        event_loop.poll(&mut events, None).expect("Poll");
        for ev in &events {
            self.ready(event_loop, ev.token(), ev.readiness())
        }
    }
}

fn main() {
    env_logger::init();
    let matches = App::new("michat")
        .arg(
            Arg::with_name("bind")
                .short("l")
                .takes_value(true)
                .required(true),
        ).get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::Poll::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    service.run(&mut event_loop);
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate elf;
extern crate rbpf;

use std::process;
use clap::{Arg, App};
use std::fs::File;
use std::io::Read;

fn main() {
    let matches = App::new("ebpf-disasm")
        .version("1.0")
        .author("Jan-Erik Rediger")
        .about("Output disassembled eBPF read from an ELF file")
        .arg(Arg::with_name("section")
             .short("s")
             .long("section")
             .value_name("SECTION")
             .help("Specify a section name")
             .takes_value(true))
        .arg(Arg::with_name("INPUT")
             .help("Sets the input ELF file to use")
             .required(true)
             .index(1))
        .arg(Arg::with_name("bytecode")
            .short("b")
            .long("bytecode")
            .help("Only show raw bytecode"))
        .arg(Arg::with_name("ccode")
            .short("c")
            .long("ccode")
            .help("Only show a C compatible bpf_insn array"))
        .arg(Arg::with_name("raw")
            .short("r")
            .long("raw")
            .help("Treat input as raw bytes"))
        .arg(Arg::with_name("list-sections")
            .short("l")
            .long("list")
            .help("List sections of the object file"))
        .get_matches();

    let section = matches.value_of("section").unwrap_or(".classifier");
    let input_file = matches.value_of("INPUT").unwrap();
    let show_bytecode = matches.is_present("bytecode");
    let show_ccode = matches.is_present("ccode");
    let raw = matches.is_present("raw");
    let list_sections = matches.is_present("list-sections");

    if show_bytecode && show_ccode {
        println!("Can't show both bytecode and C code.");
        process::exit(1);
    }

    if raw && list_sections {
        println!("Can't list sections of raw bytecode.");
        process::exit(1);
    }

    let prog;
    if raw {
        let mut file = match File::open(&input_file) {
            Ok(f) => f,
            Err(e) => {
                println!("Could not read file.");
                println!("Error: {:?}", e);
                process::exit(1);
            }
        };

        let mut buf = Vec::new();
        match file.read_to_end(&mut buf) {
            Ok(_) => {},
            Err(e) => {
                println!("Could not read file.");
                println!("Error: {:?}", e);
                process::exit(1);
            }
        }
        prog = buf;
    } else {
        let file = match elf::File::open_path(&input_file) {
            Ok(f) => f,
            Err(e) => {
                println!("Could not read elf file.");
                println!("Error: {:?}", e);
                process::exit(1);
            }
        };

        if list_sections {
            println!("Sections of {}:", input_file);
            for section in file.sections {
                println!("  {}", section.shdr.name);
            }
            process::exit(0);
        }

        let text_scn = match file.get_section(&section) {
            Some(s) => s,
            None => {
                println!("Failed to lookup '{}' section.", section);
                process::exit(1);
            }
        };

        prog = text_scn.data.clone();
    }

    if show_bytecode {
        for insn in prog.chunks(8) {
            for i in insn {
                print!("0x{:>02x}, ", i);
            }
            println!("");
        }
    } else if show_ccode {
        println!("struct bpf_insn prog[] = {{");
        for insn in prog.chunks(8) {
            print!("\t{{ ");
            print!(".code = 0x{:>02x}, ", insn[0]);
            print!(".dst_reg = 0x{:>x}, ", insn[1]&0x0f);
            print!(".src_reg = 0x{:>x}, ", insn[1]>>4);
            let off = (insn[3] as u16)<<8 | (insn[2] as u16);
            print!(".off = 0x{:>04x}, ", off);
            let imm = (insn[7] as u32)<<24 |
                (insn[6] as u32)<<16 |
                (insn[5] as u32)<<8 |
                (insn[4] as u32)<<0;
            print!(".imm = 0x{:>08x}", imm);
            println!(" }},");
        }
        println!("}};");
    } else {
        for insn in rbpf::disassembler::to_insn_vec(&prog) {
            println!("{}", insn.desc.replace(" ", "\t"));
        }
    }
}
<commit_msg>Option to prefix all printed lines with line numbers<commit_after>extern crate clap;
extern crate elf;
extern crate rbpf;

use std::process;
use clap::{Arg, App};
use std::fs::File;
use std::io::Read;

fn main() {
    let matches = App::new("ebpf-disasm")
        .version("1.0")
        .author("Jan-Erik Rediger")
        .about("Output disassembled eBPF read from an ELF file")
        .arg(Arg::with_name("section")
             .short("s")
             .long("section")
             .value_name("SECTION")
             .help("Specify a section name")
             .takes_value(true))
        .arg(Arg::with_name("INPUT")
             .help("Sets the input ELF file to use")
             .required(true)
             .index(1))
        .arg(Arg::with_name("bytecode")
            .short("b")
            .long("bytecode")
            .help("Only show raw bytecode"))
        .arg(Arg::with_name("ccode")
            .short("c")
            .long("ccode")
            .help("Only show a C compatible bpf_insn array"))
        .arg(Arg::with_name("raw")
            .short("r")
            .long("raw")
            .help("Treat input as raw bytes"))
        .arg(Arg::with_name("list-sections")
            .short("l")
            .long("list")
            .help("List sections of the object file"))
        .arg(Arg::with_name("number")
            .short("n")
            .long("number")
            .help("Number all output lines"))
        .get_matches();

    let section = matches.value_of("section").unwrap_or(".classifier");
    let input_file = matches.value_of("INPUT").unwrap();
    let show_bytecode = matches.is_present("bytecode");
    let show_ccode = matches.is_present("ccode");
    let raw = matches.is_present("raw");
    let list_sections = matches.is_present("list-sections");
    let number = matches.is_present("number");

    if show_bytecode && show_ccode {
        println!("Can't show both bytecode and C code.");
        process::exit(1);
    }

    if raw && list_sections {
        println!("Can't list sections of raw bytecode.");
        process::exit(1);
    }

    let prog;
    if raw {
        let mut file = match File::open(&input_file) {
            Ok(f) => f,
            Err(e) => {
                println!("Could not read file.");
                println!("Error: {:?}", e);
                process::exit(1);
            }
        };

        let mut buf = Vec::new();
        match file.read_to_end(&mut buf) {
            Ok(_) => {},
            Err(e) => {
                println!("Could not read file.");
                println!("Error: {:?}", e);
                process::exit(1);
            }
        }
        prog = buf;
    } else {
        let file = match elf::File::open_path(&input_file) {
            Ok(f) => f,
            Err(e) => {
                println!("Could not read elf file.");
                println!("Error: {:?}", e);
                process::exit(1);
            }
        };

        if list_sections {
            println!("Sections of {}:", input_file);
            for section in file.sections {
                println!("  {}", section.shdr.name);
            }
            process::exit(0);
        }

        let text_scn = match file.get_section(&section) {
            Some(s) => s,
            None => {
                println!("Failed to lookup '{}' section.", section);
                process::exit(1);
            }
        };

        prog = text_scn.data.clone();
    }

    if show_bytecode {
        for insn in prog.chunks(8) {
            for i in insn {
                print!("0x{:>02x}, ", i);
            }
            println!("");
        }
    } else if show_ccode {
        println!("struct bpf_insn prog[] = {{");
        for insn in prog.chunks(8) {
            print!("\t{{ ");
            print!(".code = 0x{:>02x}, ", insn[0]);
            print!(".dst_reg = 0x{:>x}, ", insn[1]&0x0f);
            print!(".src_reg = 0x{:>x}, ", insn[1]>>4);
            let off = (insn[3] as u16)<<8 | (insn[2] as u16);
            print!(".off = 0x{:>04x}, ", off);
            let imm = (insn[7] as u32)<<24 |
                (insn[6] as u32)<<16 |
                (insn[5] as u32)<<8 |
                (insn[4] as u32)<<0;
            print!(".imm = 0x{:>08x}", imm);
            println!(" }},");
        }
        println!("}};");
    } else {
        for (idx, insn) in rbpf::disassembler::to_insn_vec(&prog).iter().enumerate() {
            if number {
                print!("{:>6}  ", idx);
            }
            println!("{}", insn.desc.replace(" ", "\t"));
        }
    }
}
<|endoftext|>


<commit_before>use error::Result;
use error::Error;
use ip::IpAddr;
use hyper::Client;
use hyper::header::Connection;
use std::io;
use std::io::prelude::*;
use std::result::Result as StdResult;
use std::str::FromStr;
use regex::Regex;

#[derive(Debug)]
pub enum HttpIpProviders {
    Stdin,
    SfrLaBoxFibre,
}

impl FromStr for HttpIpProviders {
    type Err = String;

    fn from_str(s: &str) -> StdResult<HttpIpProviders, String> {
        match s {
            "stdin" => Ok(HttpIpProviders::Stdin),
            "sfrlaboxfibre" => Ok(HttpIpProviders::SfrLaBoxFibre),
            _ => Err("Unknown value for IP provider".to_owned()),
        }
    }
}

pub trait MyIPAddressProvider<'a> {
    fn get_my_ip_addr(&self) -> Result<IpAddr>;
}

pub struct StdinIpProvider;

impl<'a> MyIPAddressProvider<'a> for StdinIpProvider {
    fn get_my_ip_addr(&self) -> Result<IpAddr> {
        let mut input = String::new();

        try!(io::stdin().read_line(&mut input));

        trace!("Read stdin: {}", input);

        let ip_addr = try!(IpAddr::from_str(input.trim()));
        Ok(ip_addr)
    }
}

pub struct HttpIpProvider;

impl<'a> MyIPAddressProvider<'a> for HttpIpProvider {
    fn get_my_ip_addr(&self) -> Result<IpAddr> {
        let client = Client::new();

        let mut res = try!(client.get("http://192.168.0.1/")
            .header(Connection::close())
            .send());

        let mut body = String::new();
        try!(res.read_to_string(&mut body));

        trace!("HTTP Response: {}", body);

        let regex = try!(Regex::new(r"((?:(?:0|1[\d]{0,2}|2(?:[0-4]\d?|5[0-5]?|[6-9])?|[3-9]\d?)\.){3}(?:0|1[\d]{0,2}|2(?:[0-4]\d?|5[0-5]?|[6-9])?|[3-9]\d?))"));

        regex.captures(&*body)
            .and_then(|caps| caps.at(1))
            // Convert Option to Result
            .ok_or(Error::IpNotFound)
            .and_then(|val| IpAddr::from_str(val).map_err(|e| {
                error!("IP parse error: {}", e);
                Error::IpNotFound
            }))
    }
}
<commit_msg>Rename trait<commit_after>use error::Result;
use error::Error;
use ip::IpAddr;
use hyper::Client;
use hyper::header::Connection;
use std::io;
use std::io::prelude::*;
use std::result::Result as StdResult;
use std::str::FromStr;
use regex::Regex;

#[derive(Debug)]
pub enum HttpIpProviders {
    Stdin,
    SfrLaBoxFibre,
}

impl FromStr for HttpIpProviders {
    type Err = String;

    fn from_str(s: &str) -> StdResult<HttpIpProviders, String> {
        match s {
            "stdin" => Ok(HttpIpProviders::Stdin),
            "sfrlaboxfibre" => Ok(HttpIpProviders::SfrLaBoxFibre),
            _ => Err("Unknown value for IP provider".to_owned()),
        }
    }
}

pub trait IpAddressProvider<'a> {
    fn get_my_ip_addr(&self) -> Result<IpAddr>;
}

pub struct StdinIpProvider;

impl<'a> IpAddressProvider<'a> for StdinIpProvider {
    fn get_my_ip_addr(&self) -> Result<IpAddr> {
        let mut input = String::new();

        try!(io::stdin().read_line(&mut input));

        trace!("Read stdin: {}", input);

        let ip_addr = try!(IpAddr::from_str(input.trim()));
        Ok(ip_addr)
    }
}

pub struct HttpIpProvider;

impl<'a> IpAddressProvider<'a> for HttpIpProvider {
    fn get_my_ip_addr(&self) -> Result<IpAddr> {
        let client = Client::new();

        let mut res = try!(client.get("http://192.168.0.1/")
            .header(Connection::close())
            .send());

        let mut body = String::new();
        try!(res.read_to_string(&mut body));

        trace!("HTTP Response: {}", body);

        let regex = try!(Regex::new(r"((?:(?:0|1[\d]{0,2}|2(?:[0-4]\d?|5[0-5]?|[6-9])?|[3-9]\d?)\.){3}(?:0|1[\d]{0,2}|2(?:[0-4]\d?|5[0-5]?|[6-9])?|[3-9]\d?))"));

        regex.captures(&*body)
            .and_then(|caps| caps.at(1))
            // Convert Option to Result
            .ok_or(Error::IpNotFound)
            .and_then(|val| IpAddr::from_str(val).map_err(|e| {
                error!("IP parse error: {}", e);
                Error::IpNotFound
            }))
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Implements parallel traversal over the DOM tree.
//!
//! This traversal is based on Rayon, and therefore its safety is largely
//! verified by the type system.
//!
//! The primary trickiness and fine print for the above relates to the
//! thread safety of the DOM nodes themselves. Accessing a DOM element
//! concurrently on multiple threads is actually mostly "safe", since all
//! the mutable state is protected by an AtomicRefCell, and so we'll
//! generally panic if something goes wrong. Still, we try to to enforce our
//! thread invariants at compile time whenever possible. As such, TNode and
//! TElement are not Send, so ordinary style system code cannot accidentally
//! share them with other threads. In the parallel traversal, we explicitly
//! invoke |unsafe { SendNode::new(n) }| to put nodes in containers that may
//! be sent to other threads. This occurs in only a handful of places and is
//! easy to grep for. At the time of this writing, there is no other unsafe
//! code in the parallel traversal.

#![deny(missing_docs)]

use arrayvec::ArrayVec;
use context::{StyleContext, ThreadLocalStyleContext};
use dom::{OpaqueNode, SendNode, TElement};
use itertools::Itertools;
use rayon;
use scoped_tls::ScopedTLS;
use smallvec::SmallVec;
use traversal::{DomTraversal, PerLevelTraversalData};

/// The minimum stack size for a thread in the styling pool, in kilobytes.
pub const STYLE_THREAD_STACK_SIZE_KB: usize = 128;

/// The stack margin. If we get this deep in the stack, we will skip recursive
/// optimizations to ensure that there is sufficient room for non-recursive work.
///
/// When measured with 128KB stacks and 40KB margin, we could support 53
/// levels of recursion before the limiter kicks in, on x86_64-Linux [1].
///
/// We currently use 45KB margins, because that seems to be the minimum amount
/// of head room required by an unoptimized debug build, as measured on [2]. We
/// could probably get away with a smaller margin on optimized release builds,
/// but that would be a pain, and the extra padding reduces stability risk.
///
/// [1] See Gecko bug 1376883 for more discussion on the measurements.
/// [2] layout/style/crashtests/1383981-3.html
///
pub const STACK_SAFETY_MARGIN_KB: usize = 45;

/// The maximum number of child nodes that we will process as a single unit.
///
/// Larger values will increase style sharing cache hits and general DOM
/// locality at the expense of decreased opportunities for parallelism.  There
/// are some measurements in
/// https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c11 and comments 12
/// and 13 that investigate some slightly different values for the work unit
/// size.  If the size is significantly increased, make sure to adjust the
/// condition for kicking off a new work unit in top_down_dom, because
/// otherwise we're likely to end up doing too much work serially.  For
/// example, the condition there could become some fraction of WORK_UNIT_MAX
/// instead of WORK_UNIT_MAX.
pub const WORK_UNIT_MAX: usize = 16;

/// A set of nodes, sized to the work unit. This gets copied when sent to other
/// threads, so we keep it compact.
type WorkUnit<N> = ArrayVec<[SendNode<N>; WORK_UNIT_MAX]>;

/// A callback to create our thread local context.  This needs to be
/// out of line so we don't allocate stack space for the entire struct
/// in the caller.
#[inline(never)]
fn create_thread_local_context<'scope, E, D>(
    traversal: &'scope D,
    slot: &mut Option<ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>
{
    *slot = Some(ThreadLocalStyleContext::new(traversal.shared_context()));
}

/// A parallel top-down DOM traversal.
///
/// This algorithm traverses the DOM in a breadth-first, top-down manner. The
/// goals are:
/// * Never process a child before its parent (since child style depends on
///   parent style). If this were to happen, the styling algorithm would panic.
/// * Prioritize discovering nodes as quickly as possible to maximize
///   opportunities for parallelism.  But this needs to be weighed against
///   styling cousins on a single thread to improve sharing.
/// * Style all the children of a given node (i.e. all sibling nodes) on
///   a single thread (with an upper bound to handle nodes with an
///   abnormally large number of children). This is important because we use
///   a thread-local cache to share styles between siblings.
#[inline(always)]
#[allow(unsafe_code)]
fn top_down_dom<'a, 'scope, E, D>(nodes: &'a [SendNode<E::ConcreteNode>],
                                  root: OpaqueNode,
                                  mut traversal_data: PerLevelTraversalData,
                                  scope: &'a rayon::Scope<'scope>,
                                  pool: &'scope rayon::ThreadPool,
                                  traversal: &'scope D,
                                  tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>,
{
    debug_assert!(nodes.len() <= WORK_UNIT_MAX);

    // We set this below, when we have a borrow of the thread-local-context
    // available.
    let recursion_ok;

    // Collect all the children of the elements in our work unit. This will
    // contain the combined children of up to WORK_UNIT_MAX nodes, which may
    // be numerous. As such, we store it in a large SmallVec to minimize heap-
    // spilling, and never move it.
    let mut discovered_child_nodes = SmallVec::<[SendNode<E::ConcreteNode>; 128]>::new();
    {
        // Scope the borrow of the TLS so that the borrow is dropped before
        // a potential recursive call when we pass TailCall.
        let mut tlc = tls.ensure(
            |slot: &mut Option<ThreadLocalStyleContext<E>>| create_thread_local_context(traversal, slot));

        // Check that we're not in danger of running out of stack.
        recursion_ok = !tlc.stack_limit_checker.limit_exceeded();

        let mut context = StyleContext {
            shared: traversal.shared_context(),
            thread_local: &mut *tlc,
        };

        for n in nodes {
            // If the last node we processed produced children, we may want to
            // spawn them off into a work item. We do this at the beginning of
            // the loop (rather than at the end) so that we can traverse our
            // last bits of work directly on this thread without a spawn call.
            //
            // This has the important effect of removing the allocation and
            // context-switching overhead of the parallel traversal for perfectly
            // linear regions of the DOM, i.e.:
            //
            // <russian><doll><tag><nesting></nesting></tag></doll></russian>
            //
            // which are not at all uncommon.
            //
            // There's a tension here between spawning off a work item as soon
            // as discovered_child_nodes is nonempty and waiting until we have a
            // full work item to do so.  The former optimizes for speed of
            // discovery (we'll start discovering the kids of the things in
            // "nodes" ASAP).  The latter gives us better sharing (e.g. we can
            // share between cousins much better, because we don't hand them off
            // as separate work items, which are likely to end up on separate
            // threads) and gives us a chance to just handle everything on this
            // thread for small DOM subtrees, as in the linear example above.
            //
            // There are performance and "number of ComputedValues"
            // measurements for various testcases in
            // https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c10 and
            // following.
            //
            // The worst case behavior for waiting until we have a full work
            // item is a deep tree which has WORK_UNIT_MAX "linear" branches,
            // hence WORK_UNIT_MAX elements at each level.  Such a tree would
            // end up getting processed entirely sequentially, because we would
            // process each level one at a time as a single work unit, whether
            // via our end-of-loop tail call or not.  If we kicked off a
            // traversal as soon as we discovered kids, we would instead
            // process such a tree more or less with a thread-per-branch,
            // multiplexed across our actual threadpool.
            if discovered_child_nodes.len() >= WORK_UNIT_MAX {
                let mut traversal_data_copy = traversal_data.clone();
                traversal_data_copy.current_dom_depth += 1;
                traverse_nodes(discovered_child_nodes.drain(),
                               DispatchMode::NotTailCall,
                               recursion_ok,
                               root,
                               traversal_data_copy,
                               scope,
                               pool,
                               traversal,
                               tls);
            }

            let node = **n;
            let mut children_to_process = 0isize;
            traversal.process_preorder(&traversal_data, &mut context, node, |n| {
                children_to_process += 1;
                let send_n = unsafe { SendNode::new(n) };
                discovered_child_nodes.push(send_n);
            });

            traversal.handle_postorder_traversal(&mut context, root, node,
                                                 children_to_process);
        }
    }

    // Handle whatever elements we have queued up but not kicked off traversals
    // for yet.  If any exist, we can process them (or at least one work unit's
    // worth of them) directly on this thread by passing TailCall.
    if !discovered_child_nodes.is_empty() {
        traversal_data.current_dom_depth += 1;
        traverse_nodes(discovered_child_nodes.drain(),
                       DispatchMode::TailCall,
                       recursion_ok,
                       root,
                       traversal_data,
                       scope,
                       pool,
                       traversal,
                       tls);
    }
}

/// Controls whether traverse_nodes may make a recursive call to continue
/// doing work, or whether it should always dispatch work asynchronously.
#[derive(Clone, Copy, PartialEq)]
pub enum DispatchMode {
    /// This is the last operation by the caller.
    TailCall,
    /// This is not the last operation by the caller.
    NotTailCall,
}

impl DispatchMode {
    fn is_tail_call(&self) -> bool { matches!(*self, DispatchMode::TailCall) }
}

/// Enqueues |nodes| for processing, possibly on this thread if the tail call
/// conditions are met.
#[inline]
pub fn traverse_nodes<'a, 'scope, E, D, I>(
    nodes: I,
    mode: DispatchMode,
    recursion_ok: bool,
    root: OpaqueNode,
    traversal_data: PerLevelTraversalData,
    scope: &'a rayon::Scope<'scope>,
    pool: &'scope rayon::ThreadPool,
    traversal: &'scope D,
    tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>
)
where
    E: TElement + 'scope,
    D: DomTraversal<E>,
    I: ExactSizeIterator<Item = SendNode<E::ConcreteNode>>
{
    debug_assert_ne!(nodes.len(), 0);

    // This is a tail call from the perspective of the caller. However, we only
    // want to actually dispatch the job as a tail call if there's nothing left
    // in our local queue. Otherwise we need to return to it to maintain proper
    // breadth-first ordering. We also need to take care to avoid stack
    // overflow due to excessive tail recursion. The stack overflow avoidance
    // isn't observable to content -- we're still completely correct, just not
    // using tail recursion any more. See Gecko bugs 1368302 and 1376883.
    let may_dispatch_tail = mode.is_tail_call() &&
        recursion_ok &&
        !pool.current_thread_has_pending_tasks().unwrap();

    // In the common case, our children fit within a single work unit, in which
    // case we can pass the SmallVec directly and avoid extra allocation.
    if nodes.len() <= WORK_UNIT_MAX {
        let work: WorkUnit<E::ConcreteNode> = nodes.collect();
        if may_dispatch_tail {
            top_down_dom(&work, root,
                         traversal_data, scope, pool, traversal, tls);
        } else {
            scope.spawn(move |scope| {
                let work = work;
                top_down_dom(&work, root,
                             traversal_data, scope, pool, traversal, tls);
            });
        }
    } else {
        for chunk in nodes.chunks(WORK_UNIT_MAX).into_iter() {
            let nodes: WorkUnit<E::ConcreteNode> = chunk.collect();
            let traversal_data_copy = traversal_data.clone();
            scope.spawn(move |scope| {
                let n = nodes;
                top_down_dom(&*n, root,
                             traversal_data_copy, scope, pool, traversal, tls)
            });
        }
    }
}
<commit_msg>Increase stack safety margin for stylo.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Implements parallel traversal over the DOM tree.
//!
//! This traversal is based on Rayon, and therefore its safety is largely
//! verified by the type system.
//!
//! The primary trickiness and fine print for the above relates to the
//! thread safety of the DOM nodes themselves. Accessing a DOM element
//! concurrently on multiple threads is actually mostly "safe", since all
//! the mutable state is protected by an AtomicRefCell, and so we'll
//! generally panic if something goes wrong. Still, we try to to enforce our
//! thread invariants at compile time whenever possible. As such, TNode and
//! TElement are not Send, so ordinary style system code cannot accidentally
//! share them with other threads. In the parallel traversal, we explicitly
//! invoke |unsafe { SendNode::new(n) }| to put nodes in containers that may
//! be sent to other threads. This occurs in only a handful of places and is
//! easy to grep for. At the time of this writing, there is no other unsafe
//! code in the parallel traversal.

#![deny(missing_docs)]

use arrayvec::ArrayVec;
use context::{StyleContext, ThreadLocalStyleContext};
use dom::{OpaqueNode, SendNode, TElement};
use itertools::Itertools;
use rayon;
use scoped_tls::ScopedTLS;
use smallvec::SmallVec;
use traversal::{DomTraversal, PerLevelTraversalData};

/// The minimum stack size for a thread in the styling pool, in kilobytes.
pub const STYLE_THREAD_STACK_SIZE_KB: usize = 256;

/// The stack margin. If we get this deep in the stack, we will skip recursive
/// optimizations to ensure that there is sufficient room for non-recursive work.
///
/// We allocate large safety margins because certain OS calls can use very large
/// amounts of stack space [1]. Reserving a larger-than-necessary stack costs us
/// address space, but if we keep our safety margin big, we will generally avoid
/// committing those extra pages, and only use them in edge cases that would
/// otherwise cause crashes.
///
/// When measured with 128KB stacks and 40KB margin, we could support 53
/// levels of recursion before the limiter kicks in, on x86_64-Linux [2]. When
/// we doubled the stack size, we added it all to the safety margin, so we should
/// be able to get the same amount of recursion.
///
/// [1] https://bugzilla.mozilla.org/show_bug.cgi?id=1395708#c15
/// [2] See Gecko bug 1376883 for more discussion on the measurements.
///
pub const STACK_SAFETY_MARGIN_KB: usize = 168;

/// The maximum number of child nodes that we will process as a single unit.
///
/// Larger values will increase style sharing cache hits and general DOM
/// locality at the expense of decreased opportunities for parallelism.  There
/// are some measurements in
/// https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c11 and comments 12
/// and 13 that investigate some slightly different values for the work unit
/// size.  If the size is significantly increased, make sure to adjust the
/// condition for kicking off a new work unit in top_down_dom, because
/// otherwise we're likely to end up doing too much work serially.  For
/// example, the condition there could become some fraction of WORK_UNIT_MAX
/// instead of WORK_UNIT_MAX.
pub const WORK_UNIT_MAX: usize = 16;

/// A set of nodes, sized to the work unit. This gets copied when sent to other
/// threads, so we keep it compact.
type WorkUnit<N> = ArrayVec<[SendNode<N>; WORK_UNIT_MAX]>;

/// A callback to create our thread local context.  This needs to be
/// out of line so we don't allocate stack space for the entire struct
/// in the caller.
#[inline(never)]
fn create_thread_local_context<'scope, E, D>(
    traversal: &'scope D,
    slot: &mut Option<ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>
{
    *slot = Some(ThreadLocalStyleContext::new(traversal.shared_context()));
}

/// A parallel top-down DOM traversal.
///
/// This algorithm traverses the DOM in a breadth-first, top-down manner. The
/// goals are:
/// * Never process a child before its parent (since child style depends on
///   parent style). If this were to happen, the styling algorithm would panic.
/// * Prioritize discovering nodes as quickly as possible to maximize
///   opportunities for parallelism.  But this needs to be weighed against
///   styling cousins on a single thread to improve sharing.
/// * Style all the children of a given node (i.e. all sibling nodes) on
///   a single thread (with an upper bound to handle nodes with an
///   abnormally large number of children). This is important because we use
///   a thread-local cache to share styles between siblings.
#[inline(always)]
#[allow(unsafe_code)]
fn top_down_dom<'a, 'scope, E, D>(nodes: &'a [SendNode<E::ConcreteNode>],
                                  root: OpaqueNode,
                                  mut traversal_data: PerLevelTraversalData,
                                  scope: &'a rayon::Scope<'scope>,
                                  pool: &'scope rayon::ThreadPool,
                                  traversal: &'scope D,
                                  tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>,
{
    debug_assert!(nodes.len() <= WORK_UNIT_MAX);

    // We set this below, when we have a borrow of the thread-local-context
    // available.
    let recursion_ok;

    // Collect all the children of the elements in our work unit. This will
    // contain the combined children of up to WORK_UNIT_MAX nodes, which may
    // be numerous. As such, we store it in a large SmallVec to minimize heap-
    // spilling, and never move it.
    let mut discovered_child_nodes = SmallVec::<[SendNode<E::ConcreteNode>; 128]>::new();
    {
        // Scope the borrow of the TLS so that the borrow is dropped before
        // a potential recursive call when we pass TailCall.
        let mut tlc = tls.ensure(
            |slot: &mut Option<ThreadLocalStyleContext<E>>| create_thread_local_context(traversal, slot));

        // Check that we're not in danger of running out of stack.
        recursion_ok = !tlc.stack_limit_checker.limit_exceeded();

        let mut context = StyleContext {
            shared: traversal.shared_context(),
            thread_local: &mut *tlc,
        };

        for n in nodes {
            // If the last node we processed produced children, we may want to
            // spawn them off into a work item. We do this at the beginning of
            // the loop (rather than at the end) so that we can traverse our
            // last bits of work directly on this thread without a spawn call.
            //
            // This has the important effect of removing the allocation and
            // context-switching overhead of the parallel traversal for perfectly
            // linear regions of the DOM, i.e.:
            //
            // <russian><doll><tag><nesting></nesting></tag></doll></russian>
            //
            // which are not at all uncommon.
            //
            // There's a tension here between spawning off a work item as soon
            // as discovered_child_nodes is nonempty and waiting until we have a
            // full work item to do so.  The former optimizes for speed of
            // discovery (we'll start discovering the kids of the things in
            // "nodes" ASAP).  The latter gives us better sharing (e.g. we can
            // share between cousins much better, because we don't hand them off
            // as separate work items, which are likely to end up on separate
            // threads) and gives us a chance to just handle everything on this
            // thread for small DOM subtrees, as in the linear example above.
            //
            // There are performance and "number of ComputedValues"
            // measurements for various testcases in
            // https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c10 and
            // following.
            //
            // The worst case behavior for waiting until we have a full work
            // item is a deep tree which has WORK_UNIT_MAX "linear" branches,
            // hence WORK_UNIT_MAX elements at each level.  Such a tree would
            // end up getting processed entirely sequentially, because we would
            // process each level one at a time as a single work unit, whether
            // via our end-of-loop tail call or not.  If we kicked off a
            // traversal as soon as we discovered kids, we would instead
            // process such a tree more or less with a thread-per-branch,
            // multiplexed across our actual threadpool.
            if discovered_child_nodes.len() >= WORK_UNIT_MAX {
                let mut traversal_data_copy = traversal_data.clone();
                traversal_data_copy.current_dom_depth += 1;
                traverse_nodes(discovered_child_nodes.drain(),
                               DispatchMode::NotTailCall,
                               recursion_ok,
                               root,
                               traversal_data_copy,
                               scope,
                               pool,
                               traversal,
                               tls);
            }

            let node = **n;
            let mut children_to_process = 0isize;
            traversal.process_preorder(&traversal_data, &mut context, node, |n| {
                children_to_process += 1;
                let send_n = unsafe { SendNode::new(n) };
                discovered_child_nodes.push(send_n);
            });

            traversal.handle_postorder_traversal(&mut context, root, node,
                                                 children_to_process);
        }
    }

    // Handle whatever elements we have queued up but not kicked off traversals
    // for yet.  If any exist, we can process them (or at least one work unit's
    // worth of them) directly on this thread by passing TailCall.
    if !discovered_child_nodes.is_empty() {
        traversal_data.current_dom_depth += 1;
        traverse_nodes(discovered_child_nodes.drain(),
                       DispatchMode::TailCall,
                       recursion_ok,
                       root,
                       traversal_data,
                       scope,
                       pool,
                       traversal,
                       tls);
    }
}

/// Controls whether traverse_nodes may make a recursive call to continue
/// doing work, or whether it should always dispatch work asynchronously.
#[derive(Clone, Copy, PartialEq)]
pub enum DispatchMode {
    /// This is the last operation by the caller.
    TailCall,
    /// This is not the last operation by the caller.
    NotTailCall,
}

impl DispatchMode {
    fn is_tail_call(&self) -> bool { matches!(*self, DispatchMode::TailCall) }
}

/// Enqueues |nodes| for processing, possibly on this thread if the tail call
/// conditions are met.
#[inline]
pub fn traverse_nodes<'a, 'scope, E, D, I>(
    nodes: I,
    mode: DispatchMode,
    recursion_ok: bool,
    root: OpaqueNode,
    traversal_data: PerLevelTraversalData,
    scope: &'a rayon::Scope<'scope>,
    pool: &'scope rayon::ThreadPool,
    traversal: &'scope D,
    tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>
)
where
    E: TElement + 'scope,
    D: DomTraversal<E>,
    I: ExactSizeIterator<Item = SendNode<E::ConcreteNode>>
{
    debug_assert_ne!(nodes.len(), 0);

    // This is a tail call from the perspective of the caller. However, we only
    // want to actually dispatch the job as a tail call if there's nothing left
    // in our local queue. Otherwise we need to return to it to maintain proper
    // breadth-first ordering. We also need to take care to avoid stack
    // overflow due to excessive tail recursion. The stack overflow avoidance
    // isn't observable to content -- we're still completely correct, just not
    // using tail recursion any more. See Gecko bugs 1368302 and 1376883.
    let may_dispatch_tail = mode.is_tail_call() &&
        recursion_ok &&
        !pool.current_thread_has_pending_tasks().unwrap();

    // In the common case, our children fit within a single work unit, in which
    // case we can pass the SmallVec directly and avoid extra allocation.
    if nodes.len() <= WORK_UNIT_MAX {
        let work: WorkUnit<E::ConcreteNode> = nodes.collect();
        if may_dispatch_tail {
            top_down_dom(&work, root,
                         traversal_data, scope, pool, traversal, tls);
        } else {
            scope.spawn(move |scope| {
                let work = work;
                top_down_dom(&work, root,
                             traversal_data, scope, pool, traversal, tls);
            });
        }
    } else {
        for chunk in nodes.chunks(WORK_UNIT_MAX).into_iter() {
            let nodes: WorkUnit<E::ConcreteNode> = chunk.collect();
            let traversal_data_copy = traversal_data.clone();
            scope.spawn(move |scope| {
                let n = nodes;
                top_down_dom(&*n, root,
                             traversal_data_copy, scope, pool, traversal, tls)
            });
        }
    }
}
<|endoftext|>


<commit_before>use core::cell::UnsafeCell;
use core::sync::atomic::{AtomicUsize, Ordering};

use util::cpu_relax;

/// A synchronization primitive which can be used to run a one-time global
/// initialization. Unlike its std equivalent, this is generalized so that The
/// closure returns a value and it is stored. Once therefore acts something like
/// 1a future, too.
///
/// # Examples
///
/// ```
/// #![feature(const_fn)]
/// use spin;
///
/// static START: spin::Once<()> = spin::Once::new();
///
/// START.call_once(|| {
///     // run initialization here
/// });
/// ```
pub struct Once<T> {
    state: AtomicUsize,
    data: UnsafeCell<Option<T>>, // TODO remove option and use mem::uninitialized
}

// Same unsafe impls as `std::sync::RwLock`, because this also allows for
// concurrent reads.
unsafe impl<T: Sync + Sync> Sync for Once<T> {}
unsafe impl<T: Sync + Sync> Send for Once<T> {}

// Four states that a Once can be in, encoded into the lower bits of `state` in
// the Once structure.
const INCOMPLETE: usize = 0x0;
const RUNNING: usize = 0x1;
const COMPLETE: usize = 0x2;
const PANICKED: usize = 0x3;

#[cfg(feature = "core_intrinsics")]
#[inline(always)]
fn unreachable() -> ! {
    unsafe { ::core::intrinsics::unreachable() }
}

#[cfg(not(feature = "core_intrinsics"))]
#[inline(always)]
fn unreachable() -> ! {
    unreachable!()
}

impl<T> Once<T> {
    /// Creates a new `Once` value.
    pub const fn new() -> Once<T> {
        Once {
            state: AtomicUsize::new(INCOMPLETE),
            data: UnsafeCell::new(None),
        }
    }

    fn force_get<'a>(&'a self) -> &'a T {
        match unsafe { &*self.data.get() }.as_ref() {
            None    => unreachable(),
            Some(p) => p,
        }
    }

    /// Performs an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `call_once` has been called,
    /// and otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling thread if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified). The
    /// returned pointer will point to the result from the closure that was
    /// ran.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(const_fn)]
    /// use spin;
    ///
    /// static INIT: spin::Once<usize> = spin::Once::new();
    ///
    /// fn get_cached_val() -> usize {
    ///     *INIT.call_once(expensive_computation)
    /// }
    ///
    /// fn expensive_computation() -> usize {
    ///     // ...
    /// # 2
    /// }
    /// ```
    pub fn call_once<'a, F>(&'a self, builder: F) -> &'a T
        where F: FnOnce() -> T
    {
        let mut status = self.state.load(Ordering::SeqCst);

        if status == INCOMPLETE {
            status = self.state.compare_and_swap(INCOMPLETE,
                                                 RUNNING,
                                                 Ordering::SeqCst);
            if status == INCOMPLETE { // We init
                // We use a guard (Finish) to catch panics caused by builder
                let mut finish = Finish { state: &self.state, panicked: true };
                unsafe { *self.data.get() = Some(builder()) };
                finish.panicked = false;

                status = COMPLETE;
                self.state.store(status, Ordering::SeqCst);

                // This next line is strictly an optomization
                return self.force_get();
            }
        }

        loop {
            match status {
                INCOMPLETE => unreachable!(),
                RUNNING => { // We spin
                    cpu_relax();
                    status = self.state.load(Ordering::SeqCst)
                },
                PANICKED => panic!("Once has panicked"),
                COMPLETE => return self.force_get(),
                _ => unreachable(),
            }
        }
    }

    /// Returns a pointer iff the `Once` was previously initialized
    pub fn try<'a>(&'a self) -> Option<&'a T> {
        match self.state.load(Ordering::SeqCst) {
            COMPLETE => Some(self.force_get()),
            _        => None,
        }
    }

    /// Like try, but will spin if the `Once` is in the process of being
    /// initialized
    pub fn wait<'a>(&'a self) -> Option<&'a T> {
        loop {
            match self.state.load(Ordering::SeqCst) {
                INCOMPLETE => return None,
                RUNNING    => cpu_relax(), // We spin
                COMPLETE   => return Some(self.force_get()),
                PANICKED   => panic!("Once has panicked"),
                _ => unreachable(),
            }
        }
    }
}

struct Finish<'a> {
    state: &'a AtomicUsize,
    panicked: bool,
}

impl<'a> Drop for Finish<'a> {
    fn drop(&mut self) {
        if self.panicked {
            self.state.store(PANICKED, Ordering::SeqCst);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::prelude::v1::*;

    use std::sync::mpsc::channel;
    use std::thread;
    use super::Once;

    #[test]
    fn smoke_once() {
        static O: Once<()> = Once::new();
        let mut a = 0;
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
    }

    #[test]
    fn smoke_once_value() {
        static O: Once<usize> = Once::new();
        let a = O.call_once(|| 1);
        assert_eq!(*a, 1);
        let b = O.call_once(|| 2);
        assert_eq!(*b, 1);
    }

    #[test]
    fn stampede_once() {
        static O: Once<()> = Once::new();
        static mut run: bool = false;

        let (tx, rx) = channel();
        for _ in 0..10 {
            let tx = tx.clone();
            thread::spawn(move|| {
                for _ in 0..4 { thread::yield_now() }
                unsafe {
                    O.call_once(|| {
                        assert!(!run);
                        run = true;
                    });
                    assert!(run);
                }
                tx.send(()).unwrap();
            });
        }

        unsafe {
            O.call_once(|| {
                assert!(!run);
                run = true;
            });
            assert!(run);
        }

        for _ in 0..10 {
            rx.recv().unwrap();
        }
    }

    #[test]
    fn try() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.try().is_none());
        INIT.call_once(|| 2);
        assert_eq!(INIT.try().map(|r| *r), Some(2));
    }

    #[test]
    fn try_no_wait() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.try().is_none());
        thread::spawn(move|| {
            INIT.call_once(|| loop { });
        });
        assert!(INIT.try().is_none());
    }

    #[test]
    fn wait() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.wait().is_none());
        INIT.call_once(|| 3);
        assert_eq!(INIT.wait().map(|r| *r), Some(3));
    }

    #[test]
    fn panic() {
        use ::std::panic;

        static INIT: Once<()> = Once::new();

        // poison the once
        let t = panic::catch_unwind(|| {
            INIT.call_once(|| panic!());
        });
        assert!(t.is_err());

        // poisoning propagates
        let t = panic::catch_unwind(|| {
            INIT.call_once(|| {});
        });
        assert!(t.is_err());
    }
}
<commit_msg>Wrong Sync + Sync bound in Once<commit_after>use core::cell::UnsafeCell;
use core::sync::atomic::{AtomicUsize, Ordering};

use util::cpu_relax;

/// A synchronization primitive which can be used to run a one-time global
/// initialization. Unlike its std equivalent, this is generalized so that The
/// closure returns a value and it is stored. Once therefore acts something like
/// 1a future, too.
///
/// # Examples
///
/// ```
/// #![feature(const_fn)]
/// use spin;
///
/// static START: spin::Once<()> = spin::Once::new();
///
/// START.call_once(|| {
///     // run initialization here
/// });
/// ```
pub struct Once<T> {
    state: AtomicUsize,
    data: UnsafeCell<Option<T>>, // TODO remove option and use mem::uninitialized
}

// Same unsafe impls as `std::sync::RwLock`, because this also allows for
// concurrent reads.
unsafe impl<T: Send + Sync> Sync for Once<T> {}
unsafe impl<T: Send + Sync> Send for Once<T> {}

// Four states that a Once can be in, encoded into the lower bits of `state` in
// the Once structure.
const INCOMPLETE: usize = 0x0;
const RUNNING: usize = 0x1;
const COMPLETE: usize = 0x2;
const PANICKED: usize = 0x3;

#[cfg(feature = "core_intrinsics")]
#[inline(always)]
fn unreachable() -> ! {
    unsafe { ::core::intrinsics::unreachable() }
}

#[cfg(not(feature = "core_intrinsics"))]
#[inline(always)]
fn unreachable() -> ! {
    unreachable!()
}

impl<T> Once<T> {
    /// Creates a new `Once` value.
    pub const fn new() -> Once<T> {
        Once {
            state: AtomicUsize::new(INCOMPLETE),
            data: UnsafeCell::new(None),
        }
    }

    fn force_get<'a>(&'a self) -> &'a T {
        match unsafe { &*self.data.get() }.as_ref() {
            None    => unreachable(),
            Some(p) => p,
        }
    }

    /// Performs an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `call_once` has been called,
    /// and otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling thread if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified). The
    /// returned pointer will point to the result from the closure that was
    /// ran.
    ///
    /// # Examples
    ///
    /// ```
    /// #![feature(const_fn)]
    /// use spin;
    ///
    /// static INIT: spin::Once<usize> = spin::Once::new();
    ///
    /// fn get_cached_val() -> usize {
    ///     *INIT.call_once(expensive_computation)
    /// }
    ///
    /// fn expensive_computation() -> usize {
    ///     // ...
    /// # 2
    /// }
    /// ```
    pub fn call_once<'a, F>(&'a self, builder: F) -> &'a T
        where F: FnOnce() -> T
    {
        let mut status = self.state.load(Ordering::SeqCst);

        if status == INCOMPLETE {
            status = self.state.compare_and_swap(INCOMPLETE,
                                                 RUNNING,
                                                 Ordering::SeqCst);
            if status == INCOMPLETE { // We init
                // We use a guard (Finish) to catch panics caused by builder
                let mut finish = Finish { state: &self.state, panicked: true };
                unsafe { *self.data.get() = Some(builder()) };
                finish.panicked = false;

                status = COMPLETE;
                self.state.store(status, Ordering::SeqCst);

                // This next line is strictly an optomization
                return self.force_get();
            }
        }

        loop {
            match status {
                INCOMPLETE => unreachable!(),
                RUNNING => { // We spin
                    cpu_relax();
                    status = self.state.load(Ordering::SeqCst)
                },
                PANICKED => panic!("Once has panicked"),
                COMPLETE => return self.force_get(),
                _ => unreachable(),
            }
        }
    }

    /// Returns a pointer iff the `Once` was previously initialized
    pub fn try<'a>(&'a self) -> Option<&'a T> {
        match self.state.load(Ordering::SeqCst) {
            COMPLETE => Some(self.force_get()),
            _        => None,
        }
    }

    /// Like try, but will spin if the `Once` is in the process of being
    /// initialized
    pub fn wait<'a>(&'a self) -> Option<&'a T> {
        loop {
            match self.state.load(Ordering::SeqCst) {
                INCOMPLETE => return None,
                RUNNING    => cpu_relax(), // We spin
                COMPLETE   => return Some(self.force_get()),
                PANICKED   => panic!("Once has panicked"),
                _ => unreachable(),
            }
        }
    }
}

struct Finish<'a> {
    state: &'a AtomicUsize,
    panicked: bool,
}

impl<'a> Drop for Finish<'a> {
    fn drop(&mut self) {
        if self.panicked {
            self.state.store(PANICKED, Ordering::SeqCst);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::prelude::v1::*;

    use std::sync::mpsc::channel;
    use std::thread;
    use super::Once;

    #[test]
    fn smoke_once() {
        static O: Once<()> = Once::new();
        let mut a = 0;
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
    }

    #[test]
    fn smoke_once_value() {
        static O: Once<usize> = Once::new();
        let a = O.call_once(|| 1);
        assert_eq!(*a, 1);
        let b = O.call_once(|| 2);
        assert_eq!(*b, 1);
    }

    #[test]
    fn stampede_once() {
        static O: Once<()> = Once::new();
        static mut run: bool = false;

        let (tx, rx) = channel();
        for _ in 0..10 {
            let tx = tx.clone();
            thread::spawn(move|| {
                for _ in 0..4 { thread::yield_now() }
                unsafe {
                    O.call_once(|| {
                        assert!(!run);
                        run = true;
                    });
                    assert!(run);
                }
                tx.send(()).unwrap();
            });
        }

        unsafe {
            O.call_once(|| {
                assert!(!run);
                run = true;
            });
            assert!(run);
        }

        for _ in 0..10 {
            rx.recv().unwrap();
        }
    }

    #[test]
    fn try() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.try().is_none());
        INIT.call_once(|| 2);
        assert_eq!(INIT.try().map(|r| *r), Some(2));
    }

    #[test]
    fn try_no_wait() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.try().is_none());
        thread::spawn(move|| {
            INIT.call_once(|| loop { });
        });
        assert!(INIT.try().is_none());
    }

    #[test]
    fn wait() {
        static INIT: Once<usize> = Once::new();

        assert!(INIT.wait().is_none());
        INIT.call_once(|| 3);
        assert_eq!(INIT.wait().map(|r| *r), Some(3));
    }

    #[test]
    fn panic() {
        use ::std::panic;

        static INIT: Once<()> = Once::new();

        // poison the once
        let t = panic::catch_unwind(|| {
            INIT.call_once(|| panic!());
        });
        assert!(t.is_err());

        // poisoning propagates
        let t = panic::catch_unwind(|| {
            INIT.call_once(|| {});
        });
        assert!(t.is_err());
    }
}
<|endoftext|>


<commit_before>/// The Toml Read extensions

use toml::Value;

use tokenizer::tokenize_with_seperator;
use error::*;

pub trait TomlValueReadExt<'doc> {

    /// Extension function for reading a value from the current toml::Value document
    /// using a custom seperator
    fn read_with_seperator(&'doc self, query: &String, sep: char) -> Result<&'doc Value>;

    /// Extension function for reading a value from the current toml::Value document mutably
    /// using a custom seperator
    fn read_mut_with_seperator(&'doc mut self, query: &String, sep: char) -> Result<&'doc mut Value>;

    /// Extension function for reading a value from the current toml::Value document
    fn read(&'doc self, query: &String) -> Result<&'doc Value> {
        self.read_with_seperator(query, '.')
    }

    /// Extension function for reading a value from the current toml::Value document mutably
    fn read_mut(&'doc mut self, query: &String) -> Result<&'doc mut Value> {
        self.read_mut_with_seperator(query, '.')
    }

}

impl<'doc> TomlValueReadExt<'doc> for Value {

    fn read_with_seperator(&'doc self, query: &String, sep: char) -> Result<&'doc Value> {
        use resolver::non_mut_resolver::resolve;

        tokenize_with_seperator(query, sep).and_then(move |tokens| resolve(self, &tokens))
    }

    fn read_mut_with_seperator(&'doc mut self, query: &String, sep: char) -> Result<&'doc mut Value> {
        use resolver::mut_resolver::resolve;

        tokenize_with_seperator(query, sep).and_then(move |tokens| resolve(self, &tokens))
    }

}

<commit_msg>Add tests for Value::read() interface(s)<commit_after>/// The Toml Read extensions

use toml::Value;

use tokenizer::tokenize_with_seperator;
use error::*;

pub trait TomlValueReadExt<'doc> {

    /// Extension function for reading a value from the current toml::Value document
    /// using a custom seperator
    fn read_with_seperator(&'doc self, query: &String, sep: char) -> Result<&'doc Value>;

    /// Extension function for reading a value from the current toml::Value document mutably
    /// using a custom seperator
    fn read_mut_with_seperator(&'doc mut self, query: &String, sep: char) -> Result<&'doc mut Value>;

    /// Extension function for reading a value from the current toml::Value document
    fn read(&'doc self, query: &String) -> Result<&'doc Value> {
        self.read_with_seperator(query, '.')
    }

    /// Extension function for reading a value from the current toml::Value document mutably
    fn read_mut(&'doc mut self, query: &String) -> Result<&'doc mut Value> {
        self.read_mut_with_seperator(query, '.')
    }

}

impl<'doc> TomlValueReadExt<'doc> for Value {

    fn read_with_seperator(&'doc self, query: &String, sep: char) -> Result<&'doc Value> {
        use resolver::non_mut_resolver::resolve;

        tokenize_with_seperator(query, sep).and_then(move |tokens| resolve(self, &tokens))
    }

    fn read_mut_with_seperator(&'doc mut self, query: &String, sep: char) -> Result<&'doc mut Value> {
        use resolver::mut_resolver::resolve;

        tokenize_with_seperator(query, sep).and_then(move |tokens| resolve(self, &tokens))
    }

}

#[cfg(test)]
mod test {
    use super::*;
    use error::*;
    use toml::from_str as toml_from_str;

    #[test]
    fn test_read_empty() {
        let toml : Value = toml_from_str("").unwrap();

        let val  = toml.read_with_seperator(&String::from("a"), '.');
        assert!(val.is_err());
        let err = val.unwrap_err();

        assert!(is_match!(err.kind(), &ErrorKind::IdentifierNotFoundInDocument(_)));
    }

    #[test]
    fn test_read_table() {
        let toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let val  = toml.read_with_seperator(&String::from("table"), '.');
        assert!(val.is_ok());
        let val = val.unwrap();

        assert!(is_match!(val, &Value::Table(_)));
        match val {
            &Value::Table(ref t) => assert!(t.is_empty()),
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_read_table_value() {
        let toml : Value = toml_from_str(r#"
        [table]
        a = 1
        "#).unwrap();

        let val  = toml.read_with_seperator(&String::from("table.a"), '.');
        assert!(val.is_ok());
        let val = val.unwrap();

        assert!(is_match!(val, &Value::Integer(1)));
    }

    #[test]
    fn test_read_empty_table_value() {
        let toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let val  = toml.read_with_seperator(&String::from("table.a"), '.');
        assert!(val.is_err());
        let err = val.unwrap_err();

        assert!(is_match!(err.kind(), &ErrorKind::IdentifierNotFoundInDocument(_)));
    }

    #[test]
    fn test_read_table_index() {
        let toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let val  = toml.read_with_seperator(&String::from("table.[0]"), '.');
        assert!(val.is_err());
        let err = val.unwrap_err();

        assert!(is_match!(err.kind(), &ErrorKind::NoIndexInTable(_)));
    }

    ///
    ///
    /// Querying without specifying the seperator
    ///
    ///

    #[test]
    fn test_read_empty_without_seperator() {
        let toml : Value = toml_from_str("").unwrap();

        let val  = toml.read(&String::from("a"));
        assert!(val.is_err());
        let err = val.unwrap_err();

        assert!(is_match!(err.kind(), &ErrorKind::IdentifierNotFoundInDocument(_)));
    }

    #[test]
    fn test_read_table_without_seperator() {
        let toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let val  = toml.read(&String::from("table"));
        assert!(val.is_ok());
        let val = val.unwrap();

        assert!(is_match!(val, &Value::Table(_)));
        match val {
            &Value::Table(ref t) => assert!(t.is_empty()),
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_read_table_value_without_seperator() {
        let toml : Value = toml_from_str(r#"
        [table]
        a = 1
        "#).unwrap();

        let val  = toml.read(&String::from("table.a"));
        assert!(val.is_ok());
        let val = val.unwrap();

        assert!(is_match!(val, &Value::Integer(1)));
    }

    #[test]
    fn test_read_empty_table_value_without_seperator() {
        let toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let val  = toml.read(&String::from("table.a"));
        assert!(val.is_err());
        let err = val.unwrap_err();

        assert!(is_match!(err.kind(), &ErrorKind::IdentifierNotFoundInDocument(_)));
    }

    #[test]
    fn test_read_table_index_without_seperator() {
        let toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let val  = toml.read(&String::from("table.[0]"));
        assert!(val.is_err());
        let err = val.unwrap_err();

        assert!(is_match!(err.kind(), &ErrorKind::NoIndexInTable(_)));
    }

}

<|endoftext|>


<commit_before>use crate::lexer;
use crate::parser;
use crate::evaluation::evaluator;
use crate::evaluation::object::ObjectT;
use std::io::{stdin, stdout, Write};
use std::collections::HashMap;

const PROMPT: &str = "clojurium $ ";

pub fn start() {
    loop {
        let mut user_input = String::new();
        print!("{}", PROMPT);

        let _ = stdout().flush();

        stdin()
            .read_line(&mut user_input)
            .expect("Fatal error: Cannot read user input");

        // Something like trim (perhaps just use `trim`?)
        trim_input(&mut user_input);

        // HANDLE SPECIAL REPL CODES
        match user_input.as_str() {
            ":q" => {
                println!("Bye! Have a nice day!");
                break;
            }
            _ => 1, // just ignore this for now
        };

        let lexer = lexer::Lexer::new(user_input);
        let mut parser = parser::Parser::new(lexer);

        // TODO: Very ugly interface to parser.
        // LambdaParsers is a hack itself, so worth
        // to change it.
        let mut lambda_parsers = parser::LambdaParsers {
            prefix_parse_fns: HashMap::new(),
            infix_parse_fns: HashMap::new(),
        };

        lambda_parsers.register_parsers();

        let program = parser.parse_program(&lambda_parsers);

        // We would like to accumulate every error in program
        // and later render them to user.

        // This function used extensively in parser tests
        if !parser.errors.is_empty() {
            println!("Parser encountered {} errors", parser.errors.len());

            for error in parser.errors {
                println!("parser error: {}", error);
            }
        } else {
            let evaluated = evaluator::eval(evaluator::WN::P(program));
            println!("{}", evaluated.inspect())
        }
    }
}

fn trim_input(user_input: &mut String) {
    // Something like trim (perhaps just use `trim`?)
    if let Some('\n') = user_input.chars().next_back() {
        user_input.pop();
    }
    if let Some('\r') = user_input.chars().next_back() {
        user_input.pop();
    }
}
<commit_msg>indentation<commit_after>use crate::lexer;
use crate::parser;
use crate::evaluation::evaluator;
use crate::evaluation::object::ObjectT;
use std::io::{stdin, stdout, Write};
use std::collections::HashMap;

const PROMPT: &str = "clojurium $ ";

pub fn start() {
    loop {
        let mut user_input = String::new();
        print!("{}", PROMPT);

        let _ = stdout().flush();

        stdin()
            .read_line(&mut user_input)
            .expect("Fatal error: Cannot read user input");

        // Something like trim (perhaps just use `trim`?)
        trim_input(&mut user_input);

        // HANDLE SPECIAL REPL CODES
        match user_input.as_str() {
            ":q" => {
                println!("Bye! Have a nice day!");
                break;
            }
            _ => 1, // just ignore this for now
        };

        let lexer = lexer::Lexer::new(user_input);
        let mut parser = parser::Parser::new(lexer);

        // TODO: Very ugly interface to parser.
        // LambdaParsers is a hack itself, so worth
        // to change it.
        let mut lambda_parsers = parser::LambdaParsers {
            prefix_parse_fns: HashMap::new(),
            infix_parse_fns: HashMap::new(),
        };

        lambda_parsers.register_parsers();

        let program = parser.parse_program(&lambda_parsers);

        // We would like to accumulate every error in program
        // and later render them to user.

        // This function used extensively in parser tests
        if !parser.errors.is_empty() {
            println!("Parser encountered {} errors", parser.errors.len());

            for error in parser.errors {
                println!("parser error: {}", error);
            }
        } else {
            let evaluated = evaluator::eval(evaluator::WN::P(program));
            println!("{}", evaluated.inspect())
        }
    }
}

fn trim_input(user_input: &mut String) {
    // Something like trim (perhaps just use `trim`?)
    if let Some('\n') = user_input.chars().next_back() {
        user_input.pop();
    }
    if let Some('\r') = user_input.chars().next_back() {
        user_input.pop();
    }
}
<|endoftext|>