<commit_before>use cmp;
use io::{self, Error, ErrorKind, Result, IoVec, IoVecMut};
use mem;
use net::{SocketAddr, Shutdown};
use path::Path;
use sys::fs::{File, OpenOptions};
use sys::syscall::TimeSpec;
use sys_common::{AsInner, FromInner, IntoInner};
use time::Duration;

use super::{path_to_peer_addr, path_to_local_addr};

#[derive(Debug)]
pub struct TcpStream(File);

impl TcpStream {
    pub fn connect(addr: Result<&SocketAddr>) -> Result<TcpStream> {
        let path = format!("tcp:{}", addr?);
        let mut options = OpenOptions::new();
        options.read(true);
        options.write(true);
        Ok(TcpStream(File::open(Path::new(path.as_str()), &options)?))
    }

    pub fn connect_timeout(_addr: &SocketAddr, _timeout: Duration) -> Result<TcpStream> {
        Err(Error::new(ErrorKind::Other, "TcpStream::connect_timeout not implemented"))
    }

    pub fn duplicate(&self) -> Result<TcpStream> {
        Ok(TcpStream(self.0.dup(&[])?))
    }

    pub fn read(&self, buf: &mut [u8]) -> Result<usize> {
        self.0.read(buf)
    }

    pub fn read_vectored(&self, buf: &mut [IoVecMut<'_>]) -> io::Result<usize> {
        match buf.iter_mut().find(|b| !b.is_empty()) {
            Some(buf) => self.read(buf),
            None => Ok(0),
        }
    }

    pub fn write(&self, buf: &[u8]) -> Result<usize> {
        self.0.write(buf)
    }

    pub fn write_vectored(&self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
        match buf.iter().find(|b| !b.is_empty()) {
            Some(buf) => self.write(buf),
            None => Ok(0),
        }
    }

    pub fn take_error(&self) -> Result<Option<Error>> {
        Ok(None)
    }

    pub fn peer_addr(&self) -> Result<SocketAddr> {
        let path = self.0.path()?;
        Ok(path_to_peer_addr(path.to_str().unwrap_or("")))
    }

    pub fn socket_addr(&self) -> Result<SocketAddr> {
        let path = self.0.path()?;
        Ok(path_to_local_addr(path.to_str().unwrap_or("")))
    }

    pub fn peek(&self, _buf: &mut [u8]) -> Result<usize> {
        Err(Error::new(ErrorKind::Other, "TcpStream::peek not implemented"))
    }

    pub fn shutdown(&self, _how: Shutdown) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpStream::shutdown not implemented"))
    }

    pub fn nodelay(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpStream::nodelay not implemented"))
    }

    pub fn nonblocking(&self) -> Result<bool> {
        self.0.fd().nonblocking()
    }

    pub fn only_v6(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpStream::only_v6 not implemented"))
    }

    pub fn ttl(&self) -> Result<u32> {
        let mut ttl = [0];
        let file = self.0.dup(b"ttl")?;
        file.read(&mut ttl)?;
        Ok(ttl[0] as u32)
    }

    pub fn read_timeout(&self) -> Result<Option<Duration>> {
        let mut time = TimeSpec::default();
        let file = self.0.dup(b"read_timeout")?;
        if file.read(&mut time)? >= mem::size_of::<TimeSpec>() {
            Ok(Some(Duration::new(time.tv_sec as u64, time.tv_nsec as u32)))
        } else {
            Ok(None)
        }
    }

    pub fn write_timeout(&self) -> Result<Option<Duration>> {
        let mut time = TimeSpec::default();
        let file = self.0.dup(b"write_timeout")?;
        if file.read(&mut time)? >= mem::size_of::<TimeSpec>() {
            Ok(Some(Duration::new(time.tv_sec as u64, time.tv_nsec as u32)))
        } else {
            Ok(None)
        }
    }

    pub fn set_nodelay(&self, _nodelay: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpStream::set_nodelay not implemented"))
    }

    pub fn set_nonblocking(&self, nonblocking: bool) -> Result<()> {
        self.0.fd().set_nonblocking(nonblocking)
    }

    pub fn set_only_v6(&self, _only_v6: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpStream::set_only_v6 not implemented"))
    }

    pub fn set_ttl(&self, ttl: u32) -> Result<()> {
        let file = self.0.dup(b"ttl")?;
        file.write(&[cmp::min(ttl, 255) as u8])?;
        Ok(())
    }

    pub fn set_read_timeout(&self, duration_option: Option<Duration>) -> Result<()> {
        let file = self.0.dup(b"read_timeout")?;
        if let Some(duration) = duration_option {
            if duration.as_secs() == 0 && duration.subsec_nanos() == 0 {
                return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                          "cannot set a 0 duration timeout"));
            }
            file.write(&TimeSpec {
                tv_sec: duration.as_secs() as i64,
                tv_nsec: duration.subsec_nanos() as i32
            })?;
        } else {
            file.write(&[])?;
        }
        Ok(())
    }

    pub fn set_write_timeout(&self, duration_option: Option<Duration>) -> Result<()> {
        let file = self.0.dup(b"write_timeout")?;
        if let Some(duration) = duration_option {
            if duration.as_secs() == 0 && duration.subsec_nanos() == 0 {
                return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                          "cannot set a 0 duration timeout"));
            }
            file.write(&TimeSpec {
                tv_sec: duration.as_secs() as i64,
                tv_nsec: duration.subsec_nanos() as i32
            })?;
        } else {
            file.write(&[])?;
        }
        Ok(())
    }
}

impl AsInner<File> for TcpStream {
    fn as_inner(&self) -> &File { &self.0 }
}

impl FromInner<File> for TcpStream {
    fn from_inner(file: File) -> TcpStream {
        TcpStream(file)
    }
}

impl IntoInner<File> for TcpStream {
    fn into_inner(self) -> File { self.0 }
}

#[derive(Debug)]
pub struct TcpListener(File);

impl TcpListener {
    pub fn bind(addr: Result<&SocketAddr>) -> Result<TcpListener> {
        let path = format!("tcp:/{}", addr?);
        let mut options = OpenOptions::new();
        options.read(true);
        options.write(true);
        Ok(TcpListener(File::open(Path::new(path.as_str()), &options)?))
    }

    pub fn accept(&self) -> Result<(TcpStream, SocketAddr)> {
        let file = self.0.dup(b"listen")?;
        let path = file.path()?;
        let peer_addr = path_to_peer_addr(path.to_str().unwrap_or(""));
        Ok((TcpStream(file), peer_addr))
    }

    pub fn duplicate(&self) -> Result<TcpListener> {
        Ok(TcpListener(self.0.dup(&[])?))
    }

    pub fn take_error(&self) -> Result<Option<Error>> {
        Ok(None)
    }

    pub fn socket_addr(&self) -> Result<SocketAddr> {
        let path = self.0.path()?;
        Ok(path_to_local_addr(path.to_str().unwrap_or("")))
    }

    pub fn nonblocking(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpListener::nonblocking not implemented"))
    }

    pub fn only_v6(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpListener::only_v6 not implemented"))
    }

    pub fn ttl(&self) -> Result<u32> {
        let mut ttl = [0];
        let file = self.0.dup(b"ttl")?;
        file.read(&mut ttl)?;
        Ok(ttl[0] as u32)
    }

    pub fn set_nonblocking(&self, _nonblocking: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpListener::set_nonblocking not implemented"))
    }

    pub fn set_only_v6(&self, _only_v6: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpListener::set_only_v6 not implemented"))
    }

    pub fn set_ttl(&self, ttl: u32) -> Result<()> {
        let file = self.0.dup(b"ttl")?;
        file.write(&[cmp::min(ttl, 255) as u8])?;
        Ok(())
    }
}

impl AsInner<File> for TcpListener {
    fn as_inner(&self) -> &File { &self.0 }
}

impl FromInner<File> for TcpListener {
    fn from_inner(file: File) -> TcpListener {
        TcpListener(file)
    }
}

impl IntoInner<File> for TcpListener {
    fn into_inner(self) -> File { self.0 }
}
<commit_msg>Fix redox<commit_after>use cmp;
use io::{self, Error, ErrorKind, Result, IoVec, IoVecMut};
use mem;
use net::{SocketAddr, Shutdown};
use path::Path;
use sys::fs::{File, OpenOptions};
use sys::syscall::TimeSpec;
use sys_common::{AsInner, FromInner, IntoInner};
use time::Duration;

use super::{path_to_peer_addr, path_to_local_addr};

#[derive(Debug)]
pub struct TcpStream(File);

impl TcpStream {
    pub fn connect(addr: Result<&SocketAddr>) -> Result<TcpStream> {
        let path = format!("tcp:{}", addr?);
        let mut options = OpenOptions::new();
        options.read(true);
        options.write(true);
        Ok(TcpStream(File::open(Path::new(path.as_str()), &options)?))
    }

    pub fn connect_timeout(_addr: &SocketAddr, _timeout: Duration) -> Result<TcpStream> {
        Err(Error::new(ErrorKind::Other, "TcpStream::connect_timeout not implemented"))
    }

    pub fn duplicate(&self) -> Result<TcpStream> {
        Ok(TcpStream(self.0.dup(&[])?))
    }

    pub fn read(&self, buf: &mut [u8]) -> Result<usize> {
        self.0.read(buf)
    }

    pub fn read_vectored(&self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
        match bufs.iter_mut().find(|b| !b.is_empty()) {
            Some(buf) => self.read(buf),
            None => Ok(0),
        }
    }

    pub fn write(&self, buf: &[u8]) -> Result<usize> {
        self.0.write(buf)
    }

    pub fn write_vectored(&self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
        match bufs.iter().find(|b| !b.is_empty()) {
            Some(buf) => self.write(buf),
            None => Ok(0),
        }
    }

    pub fn take_error(&self) -> Result<Option<Error>> {
        Ok(None)
    }

    pub fn peer_addr(&self) -> Result<SocketAddr> {
        let path = self.0.path()?;
        Ok(path_to_peer_addr(path.to_str().unwrap_or("")))
    }

    pub fn socket_addr(&self) -> Result<SocketAddr> {
        let path = self.0.path()?;
        Ok(path_to_local_addr(path.to_str().unwrap_or("")))
    }

    pub fn peek(&self, _buf: &mut [u8]) -> Result<usize> {
        Err(Error::new(ErrorKind::Other, "TcpStream::peek not implemented"))
    }

    pub fn shutdown(&self, _how: Shutdown) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpStream::shutdown not implemented"))
    }

    pub fn nodelay(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpStream::nodelay not implemented"))
    }

    pub fn nonblocking(&self) -> Result<bool> {
        self.0.fd().nonblocking()
    }

    pub fn only_v6(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpStream::only_v6 not implemented"))
    }

    pub fn ttl(&self) -> Result<u32> {
        let mut ttl = [0];
        let file = self.0.dup(b"ttl")?;
        file.read(&mut ttl)?;
        Ok(ttl[0] as u32)
    }

    pub fn read_timeout(&self) -> Result<Option<Duration>> {
        let mut time = TimeSpec::default();
        let file = self.0.dup(b"read_timeout")?;
        if file.read(&mut time)? >= mem::size_of::<TimeSpec>() {
            Ok(Some(Duration::new(time.tv_sec as u64, time.tv_nsec as u32)))
        } else {
            Ok(None)
        }
    }

    pub fn write_timeout(&self) -> Result<Option<Duration>> {
        let mut time = TimeSpec::default();
        let file = self.0.dup(b"write_timeout")?;
        if file.read(&mut time)? >= mem::size_of::<TimeSpec>() {
            Ok(Some(Duration::new(time.tv_sec as u64, time.tv_nsec as u32)))
        } else {
            Ok(None)
        }
    }

    pub fn set_nodelay(&self, _nodelay: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpStream::set_nodelay not implemented"))
    }

    pub fn set_nonblocking(&self, nonblocking: bool) -> Result<()> {
        self.0.fd().set_nonblocking(nonblocking)
    }

    pub fn set_only_v6(&self, _only_v6: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpStream::set_only_v6 not implemented"))
    }

    pub fn set_ttl(&self, ttl: u32) -> Result<()> {
        let file = self.0.dup(b"ttl")?;
        file.write(&[cmp::min(ttl, 255) as u8])?;
        Ok(())
    }

    pub fn set_read_timeout(&self, duration_option: Option<Duration>) -> Result<()> {
        let file = self.0.dup(b"read_timeout")?;
        if let Some(duration) = duration_option {
            if duration.as_secs() == 0 && duration.subsec_nanos() == 0 {
                return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                          "cannot set a 0 duration timeout"));
            }
            file.write(&TimeSpec {
                tv_sec: duration.as_secs() as i64,
                tv_nsec: duration.subsec_nanos() as i32
            })?;
        } else {
            file.write(&[])?;
        }
        Ok(())
    }

    pub fn set_write_timeout(&self, duration_option: Option<Duration>) -> Result<()> {
        let file = self.0.dup(b"write_timeout")?;
        if let Some(duration) = duration_option {
            if duration.as_secs() == 0 && duration.subsec_nanos() == 0 {
                return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                          "cannot set a 0 duration timeout"));
            }
            file.write(&TimeSpec {
                tv_sec: duration.as_secs() as i64,
                tv_nsec: duration.subsec_nanos() as i32
            })?;
        } else {
            file.write(&[])?;
        }
        Ok(())
    }
}

impl AsInner<File> for TcpStream {
    fn as_inner(&self) -> &File { &self.0 }
}

impl FromInner<File> for TcpStream {
    fn from_inner(file: File) -> TcpStream {
        TcpStream(file)
    }
}

impl IntoInner<File> for TcpStream {
    fn into_inner(self) -> File { self.0 }
}

#[derive(Debug)]
pub struct TcpListener(File);

impl TcpListener {
    pub fn bind(addr: Result<&SocketAddr>) -> Result<TcpListener> {
        let path = format!("tcp:/{}", addr?);
        let mut options = OpenOptions::new();
        options.read(true);
        options.write(true);
        Ok(TcpListener(File::open(Path::new(path.as_str()), &options)?))
    }

    pub fn accept(&self) -> Result<(TcpStream, SocketAddr)> {
        let file = self.0.dup(b"listen")?;
        let path = file.path()?;
        let peer_addr = path_to_peer_addr(path.to_str().unwrap_or(""));
        Ok((TcpStream(file), peer_addr))
    }

    pub fn duplicate(&self) -> Result<TcpListener> {
        Ok(TcpListener(self.0.dup(&[])?))
    }

    pub fn take_error(&self) -> Result<Option<Error>> {
        Ok(None)
    }

    pub fn socket_addr(&self) -> Result<SocketAddr> {
        let path = self.0.path()?;
        Ok(path_to_local_addr(path.to_str().unwrap_or("")))
    }

    pub fn nonblocking(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpListener::nonblocking not implemented"))
    }

    pub fn only_v6(&self) -> Result<bool> {
        Err(Error::new(ErrorKind::Other, "TcpListener::only_v6 not implemented"))
    }

    pub fn ttl(&self) -> Result<u32> {
        let mut ttl = [0];
        let file = self.0.dup(b"ttl")?;
        file.read(&mut ttl)?;
        Ok(ttl[0] as u32)
    }

    pub fn set_nonblocking(&self, _nonblocking: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpListener::set_nonblocking not implemented"))
    }

    pub fn set_only_v6(&self, _only_v6: bool) -> Result<()> {
        Err(Error::new(ErrorKind::Other, "TcpListener::set_only_v6 not implemented"))
    }

    pub fn set_ttl(&self, ttl: u32) -> Result<()> {
        let file = self.0.dup(b"ttl")?;
        file.write(&[cmp::min(ttl, 255) as u8])?;
        Ok(())
    }
}

impl AsInner<File> for TcpListener {
    fn as_inner(&self) -> &File { &self.0 }
}

impl FromInner<File> for TcpListener {
    fn from_inner(file: File) -> TcpListener {
        TcpListener(file)
    }
}

impl IntoInner<File> for TcpListener {
    fn into_inner(self) -> File { self.0 }
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

pub static DEFAULT_ENTRY: &'static str = "\
---\
[imag]\
links = []\
version = \"0.4.0\"\
---";

/// Generator helper macro for mock app (for testing only)
///
/// Requires the following crates in scope:
///
/// * std
/// * libimagrt
/// * clap
///
#[macro_export]
macro_rules! make_mock_app {
    {
        app $appname:expr;
        module $module:ident;
        version $version:expr;
    } => {
        make_mock_app! {
            app $appname;
            module $module;
            version $version;
            with help "This is a mocking app";
        }
    };

    {
        app $appname:expr;
        modulename $module:ident;
        version $version:expr;
        with help $help:expr;
    }=> {
        mod $module {
            use clap::{App, ArgMatches};
            use libimagrt::spec::CliSpec;
            use libimagrt::runtime::Runtime;
            use libimagrt::error::RuntimeError;
            use libimagrt::configuration::{Configuration, InternalConfiguration};

            #[derive(Clone)]
            struct MockLinkApp<'a> {
                args: Vec<&'static str>,
                inner: App<'a, 'a>,
            }

            impl<'a> MockLinkApp<'a> {
                fn new(args: Vec<&'static str>) -> Self {
                    MockLinkApp {
                        args: args,
                        inner: ::build_ui(Runtime::get_default_cli_builder($appname, $version, $help)),
                    }
                }
            }

            impl<'a> CliSpec<'a> for MockLinkApp<'a> {
                fn name(&self) -> &str {
                    self.inner.get_name()
                }

                fn matches(self) -> ArgMatches<'a> {
                    self.inner.get_matches_from(self.args)
                }
            }

            impl<'a> InternalConfiguration for MockLinkApp<'a> {
                fn enable_logging(&self) -> bool {
                    false
                }

                fn use_inmemory_fs(&self) -> bool {
                    true
                }
            }

            #[allow(unused)]
            pub fn generate_minimal_test_config() -> Option<Configuration> { ::toml::de::from_str("[store]\nimplicit-create=true")
                            .map(Configuration::with_value)
                            .ok()
            }

            #[allow(unused)]
            pub fn generate_test_runtime<'a>(mut args: Vec<&'static str>) -> Result<Runtime<'a>, RuntimeError> {
                let mut cli_args = vec![$appname, "--rtp", "/tmp"];

                cli_args.append(&mut args);

                let cli_app = MockLinkApp::new(cli_args);
                Runtime::with_configuration(cli_app, generate_minimal_test_config())
            }

            #[allow(unused)]
            pub fn reset_test_runtime<'a>(mut args: Vec<&'static str>, old_runtime: Runtime)
                -> Result<Runtime<'a>, RuntimeError>
            {
                let mut cli_args = vec![$appname, "--rtp", "/tmp"];

                cli_args.append(&mut args);

                let cli_app = MockLinkApp::new(cli_args);
                Runtime::with_configuration(cli_app, generate_minimal_test_config())
                    .map(|rt| rt.with_store(old_runtime.extract_store()))
            }
        }
    };

}

<commit_msg>Remove plain-text default entry<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

/// Generator helper macro for mock app (for testing only)
///
/// Requires the following crates in scope:
///
/// * std
/// * libimagrt
/// * clap
///
#[macro_export]
macro_rules! make_mock_app {
    {
        app $appname:expr;
        module $module:ident;
        version $version:expr;
    } => {
        make_mock_app! {
            app $appname;
            module $module;
            version $version;
            with help "This is a mocking app";
        }
    };

    {
        app $appname:expr;
        modulename $module:ident;
        version $version:expr;
        with help $help:expr;
    }=> {
        mod $module {
            use clap::{App, ArgMatches};
            use libimagrt::spec::CliSpec;
            use libimagrt::runtime::Runtime;
            use libimagrt::error::RuntimeError;
            use libimagrt::configuration::{Configuration, InternalConfiguration};

            #[derive(Clone)]
            struct MockLinkApp<'a> {
                args: Vec<&'static str>,
                inner: App<'a, 'a>,
            }

            impl<'a> MockLinkApp<'a> {
                fn new(args: Vec<&'static str>) -> Self {
                    MockLinkApp {
                        args: args,
                        inner: ::build_ui(Runtime::get_default_cli_builder($appname, $version, $help)),
                    }
                }
            }

            impl<'a> CliSpec<'a> for MockLinkApp<'a> {
                fn name(&self) -> &str {
                    self.inner.get_name()
                }

                fn matches(self) -> ArgMatches<'a> {
                    self.inner.get_matches_from(self.args)
                }
            }

            impl<'a> InternalConfiguration for MockLinkApp<'a> {
                fn enable_logging(&self) -> bool {
                    false
                }

                fn use_inmemory_fs(&self) -> bool {
                    true
                }
            }

            #[allow(unused)]
            pub fn generate_minimal_test_config() -> Option<Configuration> { ::toml::de::from_str("[store]\nimplicit-create=true")
                            .map(Configuration::with_value)
                            .ok()
            }

            #[allow(unused)]
            pub fn generate_test_runtime<'a>(mut args: Vec<&'static str>) -> Result<Runtime<'a>, RuntimeError> {
                let mut cli_args = vec![$appname, "--rtp", "/tmp"];

                cli_args.append(&mut args);

                let cli_app = MockLinkApp::new(cli_args);
                Runtime::with_configuration(cli_app, generate_minimal_test_config())
            }

            #[allow(unused)]
            pub fn reset_test_runtime<'a>(mut args: Vec<&'static str>, old_runtime: Runtime)
                -> Result<Runtime<'a>, RuntimeError>
            {
                let mut cli_args = vec![$appname, "--rtp", "/tmp"];

                cli_args.append(&mut args);

                let cli_app = MockLinkApp::new(cli_args);
                Runtime::with_configuration(cli_app, generate_minimal_test_config())
                    .map(|rt| rt.with_store(old_runtime.extract_store()))
            }
        }
    };

}

<|endoftext|>


<commit_before>// Copyright (c) 2015, The Radare Project. All rights reserved.
// See the COPYING file at the top-level directory of this distribution.
// Licensed under the BSD 3-Clause License:
// <http://opensource.org/licenses/BSD-3-Clause>
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Module that implements constant propagation on the SSA IR.
//!
//! This module implements 'Sparse Conditional Constant Propagation (SCCP)' algorithm to
//! perform the propagation. For more details, please refer:
//!    * https://www.cs.utexas.edu/~lin/cs380c/wegman.pdf.
//!

use std::collections::{HashMap, VecDeque};
use middle::ssa::ssa_traits::{SSA, SSAMod};
use middle::ssa::ssa_traits::NodeType;
use middle::ir::{MArity, MOpcode};

#[derive(Copy, Clone, Debug, PartialEq)]
enum LatticeValue {
    Top,
    Bottom,
    Const(u64),
}

macro_rules! exprval_match {
   ( $x:expr ) => {
        match $x {
           ExprVal::Const(cval) => cval,
           _ => return $x
        }
   };
}

impl LatticeValue {
    fn is_undefined(&self) -> bool {
        *self == LatticeValue::Top
    }

    fn is_overdefined(&self) -> bool {
        *self == LatticeValue::Bottom
    }

    #[allow(dead_code)]
    fn is_constant(&self) -> bool {
        !(self.is_undefined() || self.is_overdefined())
    }
}

fn meet(v1: &LatticeValue, v2: &LatticeValue) -> LatticeValue {
    // Any ^ Top    = Any
    // Any ^ Bottom = Bottom
    //   C ^ C      = C      (C = Constant)
    //   C ^ D      = Bottom (C, D = Constant and C != D).

    match *v1 {
        LatticeValue::Top => return *v2,
        LatticeValue::Bottom => return LatticeValue::Bottom,
        _ => {}
    }

    match *v2 {
        LatticeValue::Top => return *v1,
        LatticeValue::Bottom => return LatticeValue::Bottom,
        _ => {}
    }

    if *v1 != *v2 {
        return LatticeValue::Bottom;
    }

    *v1
}

pub struct Analyzer<T: SSAMod + SSA + Clone> {
    ssa_worklist: VecDeque<T::ValueRef>,
    cfg_worklist: VecDeque<T::CFEdgeRef>,
    executable: HashMap<T::CFEdgeRef, bool>,
    expr_val: HashMap<T::ValueRef, LatticeValue>,
    g: T,
}

impl<T: SSA + SSAMod + Clone> Analyzer<T> {
    pub fn new(g: &mut T) -> Analyzer<T> {
        Analyzer {
            ssa_worklist: VecDeque::new(),
            cfg_worklist: VecDeque::new(),
            executable: HashMap::new(),
            expr_val: HashMap::new(),
            g: g.clone(),
        }
    }

    pub fn dump(&self) {
        println!("{:?}", self.expr_val);
    }

    fn visit_phi(&mut self, i: &T::ValueRef) -> LatticeValue {
        let operands = self.g.get_operands(i);
        let mut phi_val = self.get_value(i);

        // If "overdefined" return it.
        if phi_val.is_overdefined() {
            return LatticeValue::Bottom;
        }

        let parent_block = self.g.get_block(i);
        for op in &operands {
            let operand_block = self.g.get_block(op);
            let op_val = self.get_value(op);

            if op_val.is_undefined() {
                continue;
            }

            //if op_val.is_overdefined() {
            //    return LatticeValue::Bottom;
            //}
            // Only operand which could be executable will be considered.
            // Even these operands are overdefined.

            let edge = self.g.find_edge(&operand_block, &parent_block);
            if !self.is_executable(&edge) && edge != self.g.invalid_edge() {
                continue;
            }
            //TODO: Not sure what to do when the edge is invalid

            phi_val = meet(&phi_val, &op_val);
        }
        phi_val
    }

    fn evaluate_control_flow(&mut self, i: &T::ValueRef) {
        assert!(self.g.is_selector(i));

        let cond_val = self.get_value(i);
        let block = self.g.selects_for(i);
        let true_branch = self.g.true_edge_of(&block);
        let false_branch = self.g.false_edge_of(&block);
        match cond_val {
            LatticeValue::Bottom => {
                self.cfg_worklist.push_back(true_branch);
                self.cfg_worklist.push_back(false_branch);
            }
            LatticeValue::Top => {
                // TODO: Not really sure what to do here.
            }
            LatticeValue::Const(cval) => {
                if cval == 1 {
                    self.cfgwl_push(&true_branch);
                } else {
                    self.cfgwl_push(&false_branch);
                }
            }
        }
    }

    fn evaluate_unary_op(&mut self, i: &T::ValueRef, opcode: MOpcode) -> LatticeValue {
        let operand = self.g.get_operands(i);
        let operand = if operand.is_empty() {
            return LatticeValue::Top;
        } else {
            operand[0]
        };

        let val = self.get_value(&operand);
        let const_val = if let LatticeValue::Const(cval) = val {
            cval
        } else {
            return val;
        };

        let val = match opcode {
            MOpcode::OpWiden(_) => {
                // Nothing to do in case of widen as the value cannot change.
                const_val
            }
            MOpcode::OpNarrow(size) => {
                // Max size is 64. Therefore, we can _never_ narrow to 64.
                assert!(size < 64);
                let mask = (1 << (size)) - 1;
                const_val & mask
            }
            MOpcode::OpNot => {
                !const_val as u64
            }
            MOpcode::OpCall => {
                return LatticeValue::Bottom;
            }
            _ => unreachable!(),
        };

        LatticeValue::Const(val)
    }

    fn evaluate_binary_op(&mut self, i: &T::ValueRef, opcode: MOpcode) -> LatticeValue {
        // Do not reason about load/stores.
        match opcode {
            MOpcode::OpLoad | MOpcode::OpStore => return LatticeValue::Bottom,
            _ => { },
        }

        let operands = self.g.get_operands(i).iter().map(|x| self.get_value(x)).collect::<Vec<_>>();

        let lhs = operands[0];
        let rhs = operands[1];

        let lhs_val = if let LatticeValue::Const(cval) = lhs {
            cval
        } else {
            return lhs;
        };
        let rhs_val = if let LatticeValue::Const(cval) = rhs {
            cval
        } else {
            return rhs;
        };

        let val = match opcode {
            MOpcode::OpAdd => {
                lhs_val + rhs_val
            }
            MOpcode::OpSub => {
                if lhs_val > rhs_val {
                    lhs_val - rhs_val
                } else {
                    rhs_val - lhs_val
                }
            }
            MOpcode::OpMul => {
                lhs_val * rhs_val
            }
            MOpcode::OpDiv => {
                lhs_val / rhs_val
            }
            MOpcode::OpMod => {
                lhs_val % rhs_val
            }
            MOpcode::OpAnd => {
                lhs_val & rhs_val
            }
            MOpcode::OpOr => {
                lhs_val | rhs_val
            }
            MOpcode::OpXor => {
                lhs_val ^ rhs_val
            }
            MOpcode::OpCmp => {
                (lhs_val == rhs_val) as u64
            }
            MOpcode::OpGt => {
                (lhs_val > rhs_val) as u64
            }
            MOpcode::OpLt => {
                (lhs_val < rhs_val) as u64
            }
            MOpcode::OpLsl => {
                lhs_val << rhs_val
            }
            MOpcode::OpLsr => {
                lhs_val >> rhs_val
            }
            _ => unreachable!(),
        };

        LatticeValue::Const(val)
    }

    fn visit_expression(&mut self, i: &T::ValueRef) -> LatticeValue {
        let expr = self.g.get_node_data(i).unwrap();
        let opcode = if let NodeType::Op(opcode) = expr.nt {
            opcode
        } else {
            panic!("Found something other than an expression!");
        };

        if let MOpcode::OpConst(v) = opcode {
            return LatticeValue::Const(v as u64);
        }

        let val = match opcode.arity() {
            MArity::Unary => self.evaluate_unary_op(i, opcode),
            MArity::Binary => self.evaluate_binary_op(i, opcode),
            _ => unimplemented!(),
        };

        // If expression is a `Selector` it means that it's value can affect the
        // control flow.
        // Hence evaluate the control flow to add edges to the cfgwl.
        // TODO: Handle the case where the selector of a block may belong to a
        // different block.
        if self.g.is_selector(i) {
            self.evaluate_control_flow(i);
        }

        val
    }

    pub fn analyze(&mut self) {

        {
            let start_node = self.g.start_node();
            let edges = self.g.edges_of(&start_node);
            for &(ref next, _) in &edges {
                self.cfgwl_push(next);
            }
        }

        while !self.ssa_worklist.is_empty() || !self.cfg_worklist.is_empty() {
            while let Some(edge) = self.cfg_worklist.pop_front() {
                if !self.is_executable(&edge) {
                    self.mark_executable(&edge);
                    let block = self.g.target_of(&edge);
                    let phis = self.g.get_phis(&block);
                    for phi in &phis {
                        let v = self.visit_phi(phi);
                        self.set_value(phi, v);
                    }

                    let visits = self.g.incoming_edges(&block).iter().fold(0, |acc, &e| {
                        if self.is_executable(&e.0) {
                            acc + 1
                        } else {
                            acc
                        }
                    });

                    // If this is the first visit to the block.
                    if visits == 1 {
                        for expr in self.g.exprs_in(&block) {
                            let val = self.visit_expression(&expr);
                            self.set_value(&expr, val);
                            for use_ in self.g.get_uses(&expr) {
                                self.ssawl_push(&use_);
                            }
                        }
                    }

                    let next_edge = self.g.next_edge_of(&block);
                    if next_edge != self.g.invalid_edge() {
                        self.cfgwl_push(&next_edge);
                    }
                }
            } // End of cfgwl

            while let Some(e) = self.ssa_worklist.pop_front() {
                let t = if self.g.is_expr(&e) {
                    let block_of = self.g.block_of(&e);
                    if self.is_block_executable(&block_of) {
                        self.visit_expression(&e)
                    } else {
                        self.get_value(&e)
                    }
                } else {
                    self.visit_phi(&e)
                };

                if t != self.get_value(&e) {
                    self.set_value(&e, t);
                    for use_ in &self.g.get_uses(&e) {
                        self.ssawl_push(use_);
                    }
                }
            } // End of ssawl
        } // End of while-loop
    }

    pub fn emit_ssa(&mut self) -> T {
        for (k, v) in &self.expr_val {
            if let LatticeValue::Const(val) = *v {
                let newnode = self.g.add_const(val);
                self.g.replace(*k, newnode);
            }
        }
        let blocks = self.g.blocks();
        let mut remove_edges = Vec::<T::CFEdgeRef>::new();
        let mut remove_blocks = Vec::<T::ActionRef>::new();
        for block in &blocks {
            let edges = self.g.edges_of(block);
            for &(ref edge, _) in &edges {
                if !self.is_executable(edge) {
                    remove_edges.push(*edge);
                }
            }
            // TODO: Make this automatic in dce.
            if !self.is_block_executable(block) {
                remove_blocks.push(*block);
            }
        }
        for edge in &remove_edges {
            self.g.remove_edge(edge);
        }
        for block in &remove_blocks {
            self.g.remove_block(*block);
        }
        self.g.clone()
    }

    /// ////////////////////////////////////////////////////////////////////////
    /// / Helper functions.
    /// ////////////////////////////////////////////////////////////////////////

    fn is_executable(&self, i: &T::CFEdgeRef) -> bool {
        self.executable.get(i).cloned().unwrap_or(false)
    }

    fn mark_executable(&mut self, i: &T::CFEdgeRef) {
        let n = self.executable.entry(*i).or_insert(false);
        *n = true;
    }

    // Determines the Initial value
    fn init_val(&self, i: &T::ValueRef) -> LatticeValue {
        let node_data = self.g.get_node_data(i).unwrap();
        match node_data.nt {
            NodeType::Op(MOpcode::OpConst(v)) => LatticeValue::Const(v),
            NodeType::Undefined => LatticeValue::Bottom,
            _ => LatticeValue::Top,
        }
    }

    fn get_value(&mut self, i: &T::ValueRef) -> LatticeValue {
        if self.expr_val.contains_key(i) {
            return self.expr_val[i];
        }

        let v = self.init_val(i);
        self.expr_val.insert(*i, v);
        v
    }

    fn set_value(&mut self, i: &T::ValueRef, v: LatticeValue) {
        let n = self.expr_val.entry(*i).or_insert(LatticeValue::Top);
        *n = v;
    }

    fn is_block_executable(&self, i: &T::ActionRef) -> bool {
        // start_node is always reachable.
        if *i == self.g.start_node() {
            return true;
        }
        let incoming = self.g.incoming_edges(i);
        for &(ref edge, _) in &incoming {
            if self.is_executable(edge) {
                return true;
            }
        }
        false
    }

    fn ssawl_push(&mut self, i: &T::ValueRef) {
        if !self.g.is_expr(i) {
            return;
        }
        let owner_block = self.g.get_block(i);
        if self.is_block_executable(&owner_block) {
            self.ssa_worklist.push_back(*i);
        }
    }

    fn cfgwl_push(&mut self, i: &T::CFEdgeRef) {
        self.cfg_worklist.push_back(*i);
    }
}

#[cfg(test)]
mod test {
    use super::{LatticeValue, meet};

    #[test]
    fn test_meet() {
        let t = LatticeValue::Top;
        let b = LatticeValue::Bottom;
        let c1 = LatticeValue::Const(1);
        let c2 = LatticeValue::Const(2);

        assert_eq!(meet(&t, &t), t);
        assert_eq!(meet(&t, &b), b);
        assert_eq!(meet(&t, &c1), c1);
        assert_eq!(meet(&t, &c2), c2);
        assert_eq!(meet(&c1, &b), b);
        assert_eq!(meet(&c2, &c1), b);
        assert_eq!(meet(&c1, &c1), c1);
    }
}
<commit_msg>Fix a bug in sccp width change<commit_after>// Copyright (c) 2015, The Radare Project. All rights reserved.
// See the COPYING file at the top-level directory of this distribution.
// Licensed under the BSD 3-Clause License:
// <http://opensource.org/licenses/BSD-3-Clause>
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Module that implements constant propagation on the SSA IR.
//!
//! This module implements 'Sparse Conditional Constant Propagation (SCCP)' algorithm to
//! perform the propagation. For more details, please refer:
//!    * https://www.cs.utexas.edu/~lin/cs380c/wegman.pdf.
//!

use std::collections::{HashMap, VecDeque};
use middle::ssa::ssa_traits::{SSA, SSAMod};
use middle::ssa::ssa_traits::{NodeType, ValueType};
use middle::ir::{MArity, MOpcode};

#[derive(Copy, Clone, Debug, PartialEq)]
enum LatticeValue {
    Top,
    Bottom,
    Const(u64),
}

macro_rules! exprval_match {
   ( $x:expr ) => {
        match $x {
           ExprVal::Const(cval) => cval,
           _ => return $x
        }
   };
}

impl LatticeValue {
    fn is_undefined(&self) -> bool {
        *self == LatticeValue::Top
    }

    fn is_overdefined(&self) -> bool {
        *self == LatticeValue::Bottom
    }

    #[allow(dead_code)]
    fn is_constant(&self) -> bool {
        !(self.is_undefined() || self.is_overdefined())
    }
}

fn meet(v1: &LatticeValue, v2: &LatticeValue) -> LatticeValue {
    // Any ^ Top    = Any
    // Any ^ Bottom = Bottom
    //   C ^ C      = C      (C = Constant)
    //   C ^ D      = Bottom (C, D = Constant and C != D).

    match *v1 {
        LatticeValue::Top => return *v2,
        LatticeValue::Bottom => return LatticeValue::Bottom,
        _ => {}
    }

    match *v2 {
        LatticeValue::Top => return *v1,
        LatticeValue::Bottom => return LatticeValue::Bottom,
        _ => {}
    }

    if *v1 != *v2 {
        return LatticeValue::Bottom;
    }

    *v1
}

pub struct Analyzer<T: SSAMod + SSA + Clone> {
    ssa_worklist: VecDeque<T::ValueRef>,
    cfg_worklist: VecDeque<T::CFEdgeRef>,
    executable: HashMap<T::CFEdgeRef, bool>,
    expr_val: HashMap<T::ValueRef, LatticeValue>,
    g: T,
}

impl<T: SSA + SSAMod + Clone> Analyzer<T> {
    pub fn new(g: &mut T) -> Analyzer<T> {
        Analyzer {
            ssa_worklist: VecDeque::new(),
            cfg_worklist: VecDeque::new(),
            executable: HashMap::new(),
            expr_val: HashMap::new(),
            g: g.clone(),
        }
    }

    pub fn dump(&self) {
        println!("{:?}", self.expr_val);
    }

    fn visit_phi(&mut self, i: &T::ValueRef) -> LatticeValue {
        let operands = self.g.get_operands(i);
        let mut phi_val = self.get_value(i);

        // If "overdefined" return it.
        if phi_val.is_overdefined() {
            return LatticeValue::Bottom;
        }

        let parent_block = self.g.get_block(i);
        for op in &operands {
            let operand_block = self.g.get_block(op);
            let op_val = self.get_value(op);

            if op_val.is_undefined() {
                continue;
            }

            //if op_val.is_overdefined() {
            //    return LatticeValue::Bottom;
            //}
            // Only operand which could be executable will be considered.
            // Even these operands are overdefined.

            let edge = self.g.find_edge(&operand_block, &parent_block);
            if !self.is_executable(&edge) && edge != self.g.invalid_edge() {
                continue;
            }
            //TODO: Not sure what to do when the edge is invalid

            phi_val = meet(&phi_val, &op_val);
        }
        phi_val
    }

    fn evaluate_control_flow(&mut self, i: &T::ValueRef) {
        assert!(self.g.is_selector(i));

        let cond_val = self.get_value(i);
        let block = self.g.selects_for(i);
        let true_branch = self.g.true_edge_of(&block);
        let false_branch = self.g.false_edge_of(&block);
        match cond_val {
            LatticeValue::Bottom => {
                self.cfg_worklist.push_back(true_branch);
                self.cfg_worklist.push_back(false_branch);
            }
            LatticeValue::Top => {
                // TODO: Not really sure what to do here.
            }
            LatticeValue::Const(cval) => {
                if cval == 1 {
                    self.cfgwl_push(&true_branch);
                } else {
                    self.cfgwl_push(&false_branch);
                }
            }
        }
    }

    fn evaluate_unary_op(&mut self, i: &T::ValueRef, opcode: MOpcode) -> LatticeValue {
        let operand = self.g.get_operands(i);
        let operand = if operand.is_empty() {
            return LatticeValue::Top;
        } else {
            operand[0]
        };

        let val = self.get_value(&operand);
        let const_val = if let LatticeValue::Const(cval) = val {
            cval
        } else {
            return val;
        };

        let val = match opcode {
            MOpcode::OpWiden(_) => {
                // Nothing to do in case of widen as the value cannot change.
                const_val
            }
            MOpcode::OpNarrow(size) => {
                // Max size is 64. Therefore, we can _never_ narrow to 64.
                assert!(size < 64);
                let mask = (1 << (size)) - 1;
                const_val & mask
            }
            MOpcode::OpNot => {
                !const_val as u64
            }
            MOpcode::OpCall => {
                return LatticeValue::Bottom;
            }
            _ => unreachable!(),
        };

        LatticeValue::Const(val)
    }

    fn evaluate_binary_op(&mut self, i: &T::ValueRef, opcode: MOpcode) -> LatticeValue {
        // Do not reason about load/stores.
        match opcode {
            MOpcode::OpLoad | MOpcode::OpStore => return LatticeValue::Bottom,
            _ => { },
        }

        let operands = self.g.get_operands(i).iter().map(|x| self.get_value(x)).collect::<Vec<_>>();

        let lhs = operands[0];
        let rhs = operands[1];

        let lhs_val = if let LatticeValue::Const(cval) = lhs {
            cval
        } else {
            return lhs;
        };
        let rhs_val = if let LatticeValue::Const(cval) = rhs {
            cval
        } else {
            return rhs;
        };

        let val = match opcode {
            MOpcode::OpAdd => {
                lhs_val + rhs_val
            }
            MOpcode::OpSub => {
                if lhs_val > rhs_val {
                    lhs_val - rhs_val
                } else {
                    rhs_val - lhs_val
                }
            }
            MOpcode::OpMul => {
                lhs_val * rhs_val
            }
            MOpcode::OpDiv => {
                lhs_val / rhs_val
            }
            MOpcode::OpMod => {
                lhs_val % rhs_val
            }
            MOpcode::OpAnd => {
                lhs_val & rhs_val
            }
            MOpcode::OpOr => {
                lhs_val | rhs_val
            }
            MOpcode::OpXor => {
                lhs_val ^ rhs_val
            }
            MOpcode::OpCmp => {
                (lhs_val == rhs_val) as u64
            }
            MOpcode::OpGt => {
                (lhs_val > rhs_val) as u64
            }
            MOpcode::OpLt => {
                (lhs_val < rhs_val) as u64
            }
            MOpcode::OpLsl => {
                lhs_val << rhs_val
            }
            MOpcode::OpLsr => {
                lhs_val >> rhs_val
            }
            _ => unreachable!(),
        };

        LatticeValue::Const(val)
    }

    fn visit_expression(&mut self, i: &T::ValueRef) -> LatticeValue {
        let expr = self.g.get_node_data(i).unwrap();
        let opcode = if let NodeType::Op(opcode) = expr.nt {
            opcode
        } else {
            panic!("Found something other than an expression!");
        };

        if let MOpcode::OpConst(v) = opcode {
            return LatticeValue::Const(v as u64);
        }

        let val = match opcode.arity() {
            MArity::Unary => self.evaluate_unary_op(i, opcode),
            MArity::Binary => self.evaluate_binary_op(i, opcode),
            _ => unimplemented!(),
        };

        // If expression is a `Selector` it means that it's value can affect the
        // control flow.
        // Hence evaluate the control flow to add edges to the cfgwl.
        // TODO: Handle the case where the selector of a block may belong to a
        // different block.
        if self.g.is_selector(i) {
            self.evaluate_control_flow(i);
        }

        val
    }

    pub fn analyze(&mut self) {

        {
            let start_node = self.g.start_node();
            let edges = self.g.edges_of(&start_node);
            for &(ref next, _) in &edges {
                self.cfgwl_push(next);
            }
        }

        while !self.ssa_worklist.is_empty() || !self.cfg_worklist.is_empty() {
            while let Some(edge) = self.cfg_worklist.pop_front() {
                if !self.is_executable(&edge) {
                    self.mark_executable(&edge);
                    let block = self.g.target_of(&edge);
                    let phis = self.g.get_phis(&block);
                    for phi in &phis {
                        let v = self.visit_phi(phi);
                        self.set_value(phi, v);
                    }

                    let visits = self.g.incoming_edges(&block).iter().fold(0, |acc, &e| {
                        if self.is_executable(&e.0) {
                            acc + 1
                        } else {
                            acc
                        }
                    });

                    // If this is the first visit to the block.
                    if visits == 1 {
                        for expr in self.g.exprs_in(&block) {
                            let val = self.visit_expression(&expr);
                            self.set_value(&expr, val);
                            for use_ in self.g.get_uses(&expr) {
                                self.ssawl_push(&use_);
                            }
                        }
                    }

                    let next_edge = self.g.next_edge_of(&block);
                    if next_edge != self.g.invalid_edge() {
                        self.cfgwl_push(&next_edge);
                    }
                }
            } // End of cfgwl

            while let Some(e) = self.ssa_worklist.pop_front() {
                let t = if self.g.is_expr(&e) {
                    let block_of = self.g.block_of(&e);
                    if self.is_block_executable(&block_of) {
                        self.visit_expression(&e)
                    } else {
                        self.get_value(&e)
                    }
                } else {
                    self.visit_phi(&e)
                };

                if t != self.get_value(&e) {
                    self.set_value(&e, t);
                    for use_ in &self.g.get_uses(&e) {
                        self.ssawl_push(use_);
                    }
                }
            } // End of ssawl
        } // End of while-loop
    }

    pub fn emit_ssa(&mut self) -> T {
        for (k, v) in &self.expr_val {
            if let LatticeValue::Const(val) = *v {
                radeco_trace!("{:?} with {:?} --> Const {:#}",
                              k, self.g.get_node_data(k), val);
                // BUG: Width may be changed just using a simple replace.
                let const_node = self.g.add_const(val);
                let ndata = self.g.get_node_data(k).unwrap();
                let ValueType::Integer{ width: w } = ndata.vt;
                let new_node = if w == 64 {
                    const_node
                } else {
                    // val should not be larger than the k node could be.
                    assert!(w < 64 && val < (1 << (w)));
                    let address = self.g.get_address(k);
                    let opcode = MOpcode::OpNarrow(w as u16);
                    let new_node = self.g.add_op(opcode, ndata.vt, None);
                    self.g.set_addr(&new_node, address);
                    self.g.op_use(new_node, 0, const_node);
                    new_node
                };
                self.g.replace(*k, new_node);
            }
        }
        let blocks = self.g.blocks();
        let mut remove_edges = Vec::<T::CFEdgeRef>::new();
        let mut remove_blocks = Vec::<T::ActionRef>::new();
        for block in &blocks {
            let edges = self.g.edges_of(block);
            for &(ref edge, _) in &edges {
                if !self.is_executable(edge) {
                    remove_edges.push(*edge);
                }
            }
            // TODO: Make this automatic in dce.
            if !self.is_block_executable(block) {
                remove_blocks.push(*block);
            }
        }
        for edge in &remove_edges {
            self.g.remove_edge(edge);
        }
        for block in &remove_blocks {
            self.g.remove_block(*block);
        }
        self.g.clone()
    }

    /// ////////////////////////////////////////////////////////////////////////
    /// / Helper functions.
    /// ////////////////////////////////////////////////////////////////////////

    fn is_executable(&self, i: &T::CFEdgeRef) -> bool {
        self.executable.get(i).cloned().unwrap_or(false)
    }

    fn mark_executable(&mut self, i: &T::CFEdgeRef) {
        let n = self.executable.entry(*i).or_insert(false);
        *n = true;
    }

    // Determines the Initial value
    fn init_val(&self, i: &T::ValueRef) -> LatticeValue {
        let node_data = self.g.get_node_data(i).unwrap();
        match node_data.nt {
            NodeType::Op(MOpcode::OpConst(v)) => LatticeValue::Const(v),
            NodeType::Undefined => LatticeValue::Bottom,
            _ => LatticeValue::Top,
        }
    }

    fn get_value(&mut self, i: &T::ValueRef) -> LatticeValue {
        if self.expr_val.contains_key(i) {
            return self.expr_val[i];
        }

        let v = self.init_val(i);
        self.expr_val.insert(*i, v);
        v
    }

    fn set_value(&mut self, i: &T::ValueRef, v: LatticeValue) {
        let n = self.expr_val.entry(*i).or_insert(LatticeValue::Top);
        *n = v;
    }

    fn is_block_executable(&self, i: &T::ActionRef) -> bool {
        // start_node is always reachable.
        if *i == self.g.start_node() {
            return true;
        }
        let incoming = self.g.incoming_edges(i);
        for &(ref edge, _) in &incoming {
            if self.is_executable(edge) {
                return true;
            }
        }
        false
    }

    fn ssawl_push(&mut self, i: &T::ValueRef) {
        if !self.g.is_expr(i) {
            return;
        }
        let owner_block = self.g.get_block(i);
        if self.is_block_executable(&owner_block) {
            self.ssa_worklist.push_back(*i);
        }
    }

    fn cfgwl_push(&mut self, i: &T::CFEdgeRef) {
        self.cfg_worklist.push_back(*i);
    }
}

#[cfg(test)]
mod test {
    use super::{LatticeValue, meet};

    #[test]
    fn test_meet() {
        let t = LatticeValue::Top;
        let b = LatticeValue::Bottom;
        let c1 = LatticeValue::Const(1);
        let c2 = LatticeValue::Const(2);

        assert_eq!(meet(&t, &t), t);
        assert_eq!(meet(&t, &b), b);
        assert_eq!(meet(&t, &c1), c1);
        assert_eq!(meet(&t, &c2), c2);
        assert_eq!(meet(&c1, &b), b);
        assert_eq!(meet(&c2, &c1), b);
        assert_eq!(meet(&c1, &c1), c1);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2018 Rene van der Meer
//
// 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.

// gpio_status.rs
//
// Retrieves the mode and logic level for each of the pins on the standard
// 40-pin GPIO header, and displays the results in an ASCII table.

use std::fmt;
use std::process::exit;

use rppal::gpio::Gpio;
use rppal::system::{DeviceInfo, Model};

enum PinType {
    Gpio(u8),
    Ground,
    Power3v3,
    Power5v,
}

impl fmt::Display for PinType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            PinType::Gpio(pin) => write!(f, "GPIO{}", pin),
            PinType::Ground => write!(f, "{:<5}", "GND"),
            PinType::Power3v3 => write!(f, "{:<5}", "3.3 V"),
            PinType::Power5v => write!(f, "{:<5}", "5 V"),
        }
    }
}

const HEADER_40: [PinType; 40] = [
    PinType::Power3v3, // Physical pin 1
    PinType::Power5v,  // Physical pin 2
    PinType::Gpio(2),  // Physical pin 3
    PinType::Power5v,  // Physical pin 4
    PinType::Gpio(3),  // Physical pin 5
    PinType::Ground,   // Physical pin 6
    PinType::Gpio(4),  // Physical pin 7
    PinType::Gpio(14), // Physical pin 8
    PinType::Ground,   // Physical pin 9
    PinType::Gpio(15), // Physical pin 10
    PinType::Gpio(17), // Physical pin 11
    PinType::Gpio(18), // Physical pin 12
    PinType::Gpio(27), // Physical pin 13
    PinType::Ground,   // Physical pin 14
    PinType::Gpio(22), // Physical pin 15
    PinType::Gpio(23), // Physical pin 16
    PinType::Power3v3, // Physical pin 17
    PinType::Gpio(24), // Physical pin 18
    PinType::Gpio(10), // Physical pin 19
    PinType::Ground,   // Physical pin 20
    PinType::Gpio(9),  // Physical pin 21
    PinType::Gpio(25), // Physical pin 22
    PinType::Gpio(11), // Physical pin 23
    PinType::Gpio(8),  // Physical pin 24
    PinType::Ground,   // Physical pin 25
    PinType::Gpio(7),  // Physical pin 26
    PinType::Gpio(0),  // Physical pin 27
    PinType::Gpio(1),  // Physical pin 28
    PinType::Gpio(5),  // Physical pin 29
    PinType::Ground,   // Physical pin 30
    PinType::Gpio(6),  // Physical pin 31
    PinType::Gpio(12), // Physical pin 32
    PinType::Gpio(13), // Physical pin 33
    PinType::Ground,   // Physical pin 34
    PinType::Gpio(19), // Physical pin 35
    PinType::Gpio(16), // Physical pin 36
    PinType::Gpio(26), // Physical pin 37
    PinType::Gpio(20), // Physical pin 38
    PinType::Ground,   // Physical pin 39
    PinType::Gpio(21), // Physical pin 40
];

fn format_pin(
    buf: &mut String,
    pin: usize,
    gpio: impl fmt::Display,
    mode: impl fmt::Display,
    level: impl fmt::Display,
) {
    if pin % 2 != 0 {
        buf.push_str(&format!(
            "| {:>4} | {:<5} | {:>1} | {:>2} |",
            gpio, mode, level, pin
        ));
    } else {
        buf.push_str(&format!(
            " {:>2} | {:>1} | {:<5} | {:>4} |\n",
            pin, level, mode, gpio
        ));
    }
}

fn print_gpio_status() {
    let gpio = Gpio::new().unwrap_or_else(|e| {
        eprintln!("Error: Can't access GPIO peripheral ({})", e);
        exit(1);
    });

    let mut buf = String::with_capacity(1600);

    buf.push_str("+------+-------+---+---------+---+-------+------+\n");
    buf.push_str("| GPIO | Mode  | L |   Pin   | L | Mode  | GPIO |\n");
    buf.push_str("+------+-------+---+----+----+---+-------+------+\n");

    for (idx, pin_type) in HEADER_40.iter().enumerate() {
        match pin_type {
            PinType::Gpio(bcm_gpio) => {
                let pin = gpio.get(*bcm_gpio).unwrap_or_else(|| {
                    eprintln!("Error: Can't access GPIO pin {}", *bcm_gpio);
                    exit(1);
                });

                format_pin(
                    &mut buf,
                    idx + 1,
                    bcm_gpio,
                    format!("{}", pin.mode()).to_uppercase(),
                    pin.read() as u8,
                );
            }
            _ => format_pin(&mut buf, idx + 1, "", pin_type, ""),
        };
    }

    buf.push_str("+------+-------+---+----+----+---+-------+------+\n");

    print!("{}", buf);
}

fn main() {
    match DeviceInfo::new()
        .unwrap_or_else(|e| {
            eprintln!("Error: Can't identify Raspberry Pi model ({})", e);
            exit(1);
        })
        .model()
    {
        Model::RaspberryPiAPlus
        | Model::RaspberryPiBPlus
        | Model::RaspberryPi2B
        | Model::RaspberryPi3APlus
        | Model::RaspberryPi3B
        | Model::RaspberryPi3BPlus
        | Model::RaspberryPiZero
        | Model::RaspberryPiZeroW => print_gpio_status(),
        model => {
            eprintln!("Error: No GPIO header information available for {}", model);
            exit(1);
        }
    }
}
<commit_msg>Add support for older models with 26-pin headers<commit_after>// Copyright (c) 2017-2018 Rene van der Meer
//
// 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.

// gpio_status.rs
//
// Retrieves the mode and logic level for each of the pins on the standard
// 40-pin GPIO header, and displays the results in an ASCII table.

use std::fmt;
use std::process::exit;

use rppal::gpio::Gpio;
use rppal::system::{DeviceInfo, Model};

enum PinType {
    Gpio(u8),
    Ground,
    Power3v3,
    Power5v,
}

impl fmt::Display for PinType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            PinType::Gpio(pin) => write!(f, "GPIO{}", pin),
            PinType::Ground => write!(f, "{:<5}", "GND"),
            PinType::Power3v3 => write!(f, "{:<5}", "3.3 V"),
            PinType::Power5v => write!(f, "{:<5}", "5 V"),
        }
    }
}

const HEADER: [PinType; 40] = [
    PinType::Power3v3, // Physical pin 1
    PinType::Power5v,  // Physical pin 2
    PinType::Gpio(2),  // Physical pin 3
    PinType::Power5v,  // Physical pin 4
    PinType::Gpio(3),  // Physical pin 5
    PinType::Ground,   // Physical pin 6
    PinType::Gpio(4),  // Physical pin 7
    PinType::Gpio(14), // Physical pin 8
    PinType::Ground,   // Physical pin 9
    PinType::Gpio(15), // Physical pin 10
    PinType::Gpio(17), // Physical pin 11
    PinType::Gpio(18), // Physical pin 12
    PinType::Gpio(27), // Physical pin 13
    PinType::Ground,   // Physical pin 14
    PinType::Gpio(22), // Physical pin 15
    PinType::Gpio(23), // Physical pin 16
    PinType::Power3v3, // Physical pin 17
    PinType::Gpio(24), // Physical pin 18
    PinType::Gpio(10), // Physical pin 19
    PinType::Ground,   // Physical pin 20
    PinType::Gpio(9),  // Physical pin 21
    PinType::Gpio(25), // Physical pin 22
    PinType::Gpio(11), // Physical pin 23
    PinType::Gpio(8),  // Physical pin 24
    PinType::Ground,   // Physical pin 25
    PinType::Gpio(7),  // Physical pin 26
    PinType::Gpio(0),  // Physical pin 27
    PinType::Gpio(1),  // Physical pin 28
    PinType::Gpio(5),  // Physical pin 29
    PinType::Ground,   // Physical pin 30
    PinType::Gpio(6),  // Physical pin 31
    PinType::Gpio(12), // Physical pin 32
    PinType::Gpio(13), // Physical pin 33
    PinType::Ground,   // Physical pin 34
    PinType::Gpio(19), // Physical pin 35
    PinType::Gpio(16), // Physical pin 36
    PinType::Gpio(26), // Physical pin 37
    PinType::Gpio(20), // Physical pin 38
    PinType::Ground,   // Physical pin 39
    PinType::Gpio(21), // Physical pin 40
];

const MAX_PINS_SHORT: usize = 26;
const MAX_PINS_LONG: usize = 40;

fn format_pin(
    buf: &mut String,
    pin: usize,
    gpio: impl fmt::Display,
    mode: impl fmt::Display,
    level: impl fmt::Display,
) {
    if pin % 2 != 0 {
        buf.push_str(&format!(
            "| {:>4} | {:<5} | {:>1} | {:>2} |",
            gpio, mode, level, pin
        ));
    } else {
        buf.push_str(&format!(
            " {:>2} | {:>1} | {:<5} | {:>4} |\n",
            pin, level, mode, gpio
        ));
    }
}

fn print_header(max: usize) {
    let gpio = Gpio::new().unwrap_or_else(|e| {
        eprintln!("Error: Can't access GPIO peripheral ({})", e);
        exit(1);
    });

    let mut buf = String::with_capacity(1600);

    buf.push_str("+------+-------+---+---------+---+-------+------+\n");
    buf.push_str("| GPIO | Mode  | L |   Pin   | L | Mode  | GPIO |\n");
    buf.push_str("+------+-------+---+----+----+---+-------+------+\n");

    for (idx, pin_type) in HEADER.iter().take(max).enumerate() {
        match pin_type {
            PinType::Gpio(bcm_gpio) => {
                let pin = gpio.get(*bcm_gpio).unwrap_or_else(|| {
                    eprintln!("Error: Can't access GPIO pin {}", *bcm_gpio);
                    exit(1);
                });

                format_pin(
                    &mut buf,
                    idx + 1,
                    bcm_gpio,
                    format!("{}", pin.mode()).to_uppercase(),
                    pin.read() as u8,
                );
            }
            _ => format_pin(&mut buf, idx + 1, "", pin_type, ""),
        };
    }

    buf.push_str("+------+-------+---+----+----+---+-------+------+\n");

    print!("{}", buf);
}

fn main() {
    match DeviceInfo::new()
        .unwrap_or_else(|e| {
            eprintln!("Error: Can't identify Raspberry Pi model ({})", e);
            exit(1);
        })
        .model()
    {
        Model::RaspberryPiA | Model::RaspberryPiBRev1 | Model::RaspberryPiBRev2 => {
            print_header(MAX_PINS_SHORT)
        }
        Model::RaspberryPiAPlus
        | Model::RaspberryPiBPlus
        | Model::RaspberryPi2B
        | Model::RaspberryPi3APlus
        | Model::RaspberryPi3B
        | Model::RaspberryPi3BPlus
        | Model::RaspberryPiZero
        | Model::RaspberryPiZeroW => print_header(MAX_PINS_LONG),
        model => {
            eprintln!("Error: No GPIO header information available for {}", model);
            exit(1);
        }
    }
}
<|endoftext|>


<commit_before>//! Hello world server.
//!
//! A simple server that accepts connections, writes "hello world\n", and closes
//! the connection.
//!
//! You can test this out by running:
//!
//!     cargo run --example hello_world
//!
//! And then in another terminal run:
//!
//!     telnet localhost 6142
//!

#![deny(warnings)]

extern crate tokio;

use tokio::io;
use tokio::net::TcpListener;
use tokio::prelude::*;

pub fn main() {
    let addr = "127.0.0.1:6142".parse().unwrap();

    // Bind a TCP listener to the socket address.
    //
    // Note that this is the Tokio TcpListener, which is fully async.
    let listener = TcpListener::bind(&addr).unwrap();

    // The server task asynchronously iterates over and processes each
    // incoming connection.
    let server = listener.incoming().for_each(|socket| {
        println!("accepted socket; addr={:?}", socket.peer_addr().unwrap());

        let connection = io::write_all(socket, "hello world\n")
            .then(|res| {
                println!("wrote message; success={:?}", res.is_ok());
                Ok(())
            });

        // Spawn a new task that processes the socket:
        tokio::spawn(connection);

        Ok(())
    })
    .map_err(|err| {
        // All tasks must have an `Error` type of `()`. This forces error
        // handling and helps avoid silencing failures.
        //
        // In our example, we are only going to log the error to STDOUT.
        println!("accept error = {:?}", err);
    });

    println!("server running on localhost:6142");

    // Start the Tokio runtime.
    //
    // The Tokio is a pre-configured "out of the box" runtime for building
    // asynchronous applications. It includes both a reactor and a task
    // scheduler. This means applications are multithreaded by default.
    //
    // This function blocks until the runtime reaches an idle state. Idle is
    // defined as all spawned tasks have completed and all I/O resources (TCP
    // sockets in our case) have been dropped.
    //
    // In our example, we have not defined a shutdown strategy, so this will
    // block until `ctrl-c` is pressed at the terminal.
    tokio::run(server);
}
<commit_msg>tokio: change hello world to new, simpler example (#690)<commit_after>//! Hello world server.
//!
//! A simple client that opens a TCP stream, writes "hello world\n", and closes
//! the connection.
//!
//! You can test this out by running:
//!
//!     ncat -l 6142
//!
//! And then in another terminal run:
//!
//!     cargo run --example hello_world

#![deny(warnings)]

extern crate tokio;

use tokio::io;
use tokio::net::TcpStream;
use tokio::prelude::*;

pub fn main() {
    let addr = "127.0.0.1:6142".parse().unwrap();

    // Open a TCP stream to the socket address.
    //
    // Note that this is the Tokio TcpStream, which is fully async.
    let client = TcpStream::connect(&addr).and_then(|stream| {
        println!("created stream");
        io::write_all(stream, "hello world\n").then(|result| {
            println!("wrote to stream; success={:?}", result.is_ok());
            Ok(())
        })
    })
    .map_err(|err| {
        // All tasks must have an `Error` type of `()`. This forces error
        // handling and helps avoid silencing failures.
        //
        // In our example, we are only going to log the error to STDOUT.
        println!("connection error = {:?}", err);
    });

    // Start the Tokio runtime.
    //
    // The Tokio is a pre-configured "out of the box" runtime for building
    // asynchronous applications. It includes both a reactor and a task
    // scheduler. This means applications are multithreaded by default.
    //
    // This function blocks until the runtime reaches an idle state. Idle is
    // defined as all spawned tasks have completed and all I/O resources (TCP
    // sockets in our case) have been dropped.
    println!("About to create the stream and write to it...");
    tokio::run(client);
    println!("Stream has been created and written to.");
}
<|endoftext|>


<commit_before>// Copyright 2014 The HTML5 for Rust Project Developers. See the
// COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate html5;

use std::io;
use std::default::Default;
use std::string::String;

use html5::sink::rcdom::{RcDom, Handle, Document, Doctype, Text, Comment, Element};
use html5::{parse, one_input};

// This is not proper HTML serialization, of course.

fn walk(indent: uint, handle: Handle) {
    let node = handle.borrow();
    // FIXME: don't allocate
    print!("{:s}", String::from_char(indent, ' '));
    match node.node {
        Document
            => println!("#Document"),

        Doctype(ref name, ref public, ref system)
            => println!("<!DOCTYPE {:s} \"{:s}\" \"{:s}\">", *name, *public, *system),

        Text(ref text)
            => println!("#text: {:s}", text.escape_default()),

        Comment(ref text)
            => println!("<!-- {:s} -->", text.escape_default()),

        Element(ref name, ref attrs) => {
            print!("<{:s}", name.as_slice());
            for attr in attrs.iter() {
                print!(" {:s}=\"{:s}\"", attr.name.name, attr.value);
            }
            println!(">");
        }
    }

    for child in node.children.iter() {
        walk(indent+4, child.clone());
    }
}

fn main() {
    let input = io::stdin().read_to_str().unwrap();
    let dom: RcDom = parse(one_input(input), Default::default());
    walk(0, dom.document);
}
<commit_msg>print-rcdom: Print parse errors<commit_after>// Copyright 2014 The HTML5 for Rust Project Developers. See the
// COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate html5;

use std::io;
use std::default::Default;
use std::string::String;

use html5::sink::rcdom::{RcDom, Handle, Document, Doctype, Text, Comment, Element};
use html5::{parse, one_input};

// This is not proper HTML serialization, of course.

fn walk(indent: uint, handle: Handle) {
    let node = handle.borrow();
    // FIXME: don't allocate
    print!("{:s}", String::from_char(indent, ' '));
    match node.node {
        Document
            => println!("#Document"),

        Doctype(ref name, ref public, ref system)
            => println!("<!DOCTYPE {:s} \"{:s}\" \"{:s}\">", *name, *public, *system),

        Text(ref text)
            => println!("#text: {:s}", text.escape_default()),

        Comment(ref text)
            => println!("<!-- {:s} -->", text.escape_default()),

        Element(ref name, ref attrs) => {
            print!("<{:s}", name.as_slice());
            for attr in attrs.iter() {
                print!(" {:s}=\"{:s}\"", attr.name.name, attr.value);
            }
            println!(">");
        }
    }

    for child in node.children.iter() {
        walk(indent+4, child.clone());
    }
}

fn main() {
    let input = io::stdin().read_to_str().unwrap();
    let dom: RcDom = parse(one_input(input), Default::default());
    walk(0, dom.document);

    if !dom.errors.is_empty() {
        println!("\nParse errors:");
        for err in dom.errors.move_iter() {
            println!("    {}", err);
        }
    }
}
<|endoftext|>


<commit_before>//! Implementation of a simple uTP client and server.
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate utp;

// A little macro to make it easier to unwrap return values or halt the program
macro_rules! iotry {
    ($e:expr) => (match $e { Ok(v) => v, Err(e) => panic!("{}", e), })
}

fn usage() {
    println!("Usage: utp [-s|-c] <address> <port>");
}

fn main() {
    use utp::UtpStream;
    use std::io::{stdin, stdout, stderr, Read, Write};

    // This example may run in either server or client mode.
    // Using an enum tends to make the code cleaner and easier to read.
    enum Mode {Server, Client}

    // Start logging
    env_logger::init().unwrap();

    // Fetch arguments
    let mut args = std::env::args();

    // Skip program name
    args.next();

    // Parse the mode argument
    let mode: Mode = match args.next() {
        Some(ref s) if &s[..] == "-s" => Mode::Server,
        Some(ref s) if &s[..] == "-c" => Mode::Client,
        _ => { usage(); return; }
    };

    // Parse the address argument or use a default if none is provided
    let addr = match (args.next(), args.next()) {
        (None, None) => "127.0.0.1:8080".to_string(),
        (Some(ip), Some(port)) => format!("{}:{}", ip, port),
        _ => { usage(); return; }
    };
    let addr: &str = &addr;

    match mode {
        Mode::Server => {
            // Create a listening stream
            let mut stream = iotry!(UtpStream::bind(addr));
            let mut writer = stdout();
            let _ = writeln!(&mut stderr(), "Serving on {}", addr);

            // Create a reasonably-sized buffer
            let mut payload = vec![0; 1024 * 1024];

            // Wait for a new connection and print the received data to stdout.
            // Reading and printing chunks like this feels more interactive than trying to read
            // everything with `read_to_end` and avoids resizing the buffer multiple times.
            loop {
                match stream.read(&mut payload) {
                    Ok(0) => break,
                    Ok(read) => iotry!(writer.write(&payload[..read])),
                    Err(e) => panic!("{}", e)
                };
            }
        }
        Mode::Client => {
            // Create a stream and try to connect to the remote address
            let mut stream = iotry!(UtpStream::connect(addr));
            let mut reader = stdin();

            // Create a reasonably sized buffer
            let mut payload = vec![0; 1024 * 1024];

            // Read from stdin and send it to the remote server.
            // Once again, reading and sending small chunks like this avoids having to read the
            // entire input (which may be endless!) before starting to send, unlike what would
            // happen if we were to use `read_to_end` on `reader`.
            loop {
                match reader.read(&mut payload) {
                    Ok(0) => break,
                    Ok(read) => iotry!(stream.write(&payload[..read])),
                    Err(e) => {
                        iotry!(stream.close());
                        panic!("{:?}", e);
                    }
                };
            }

            // Explicitly close the stream.
            iotry!(stream.close());
        }
    }
}
<commit_msg>Remove unnecessary hyphen.<commit_after>//! Implementation of a simple uTP client and server.
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate utp;

// A little macro to make it easier to unwrap return values or halt the program
macro_rules! iotry {
    ($e:expr) => (match $e { Ok(v) => v, Err(e) => panic!("{}", e), })
}

fn usage() {
    println!("Usage: utp [-s|-c] <address> <port>");
}

fn main() {
    use utp::UtpStream;
    use std::io::{stdin, stdout, stderr, Read, Write};

    // This example may run in either server or client mode.
    // Using an enum tends to make the code cleaner and easier to read.
    enum Mode {Server, Client}

    // Start logging
    env_logger::init().unwrap();

    // Fetch arguments
    let mut args = std::env::args();

    // Skip program name
    args.next();

    // Parse the mode argument
    let mode: Mode = match args.next() {
        Some(ref s) if &s[..] == "-s" => Mode::Server,
        Some(ref s) if &s[..] == "-c" => Mode::Client,
        _ => { usage(); return; }
    };

    // Parse the address argument or use a default if none is provided
    let addr = match (args.next(), args.next()) {
        (None, None) => "127.0.0.1:8080".to_string(),
        (Some(ip), Some(port)) => format!("{}:{}", ip, port),
        _ => { usage(); return; }
    };
    let addr: &str = &addr;

    match mode {
        Mode::Server => {
            // Create a listening stream
            let mut stream = iotry!(UtpStream::bind(addr));
            let mut writer = stdout();
            let _ = writeln!(&mut stderr(), "Serving on {}", addr);

            // Create a reasonably sized buffer
            let mut payload = vec![0; 1024 * 1024];

            // Wait for a new connection and print the received data to stdout.
            // Reading and printing chunks like this feels more interactive than trying to read
            // everything with `read_to_end` and avoids resizing the buffer multiple times.
            loop {
                match stream.read(&mut payload) {
                    Ok(0) => break,
                    Ok(read) => iotry!(writer.write(&payload[..read])),
                    Err(e) => panic!("{}", e)
                };
            }
        }
        Mode::Client => {
            // Create a stream and try to connect to the remote address
            let mut stream = iotry!(UtpStream::connect(addr));
            let mut reader = stdin();

            // Create a reasonably sized buffer
            let mut payload = vec![0; 1024 * 1024];

            // Read from stdin and send it to the remote server.
            // Once again, reading and sending small chunks like this avoids having to read the
            // entire input (which may be endless!) before starting to send, unlike what would
            // happen if we were to use `read_to_end` on `reader`.
            loop {
                match reader.read(&mut payload) {
                    Ok(0) => break,
                    Ok(read) => iotry!(stream.write(&payload[..read])),
                    Err(e) => {
                        iotry!(stream.close());
                        panic!("{:?}", e);
                    }
                };
            }

            // Explicitly close the stream.
            iotry!(stream.close());
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::path::PathBuf;

const PROTOS: &'static [&'static str] = &[
    "vendor/grr/grr/proto/grr_response_proto/semantic.proto",
    "vendor/grr/grr/proto/grr_response_proto/sysinfo.proto",
    "vendor/grr/grr/proto/grr_response_proto/knowledge_base.proto",
    "vendor/grr/grr/proto/grr_response_proto/jobs.proto",
    "vendor/grr/grr/proto/grr_response_proto/timeline.proto",
    "vendor/grr/grr/proto/grr_response_proto/anomaly.proto",
    "vendor/grr/grr/proto/grr_response_proto/export.proto",
    "vendor/grr/grr/proto/grr_response_proto/objects.proto",
    "vendor/grr/grr/proto/grr_response_proto/output_plugin.proto",
    "vendor/grr/grr/proto/grr_response_proto/flows.proto",
    "vendor/grr/grr/proto/grr_response_proto/sysinfo.proto",
    "vendor/grr/grr/proto/grr_response_proto/user.proto",
];

const INCLUDES: &'static [&'static str] = &[
    "vendor/grr/grr/proto",
];

fn main() {
    // `prost` generated code.
    prost_build::compile_protos(&PROTOS, &INCLUDES)
        .expect("failed to compile proto files");

    // There is a problem with one enum generated by PROST!: it's values use
    // name mangling, but it's default value does not. This is likely a bug in
    // PROST! itself, but for now we hack around it by replacing the spurious
    // line in the output file ourselves.
    let outdir: PathBuf = std::env::var("OUT_DIR")
        .expect("no output directory")
        .into();

    let target = outdir.join("grr.rs");

    let grr = std::fs::read_to_string(&target)
        .expect("failed to read the generated output file")
        .replace("TskFsAttrTypeDefault", "Default");

    std::fs::write(&target, grr)
        .expect("failed to write updated output file");

    // `protobuf` generated code.
    let proto_out_dir = outdir.join("proto");
    std::fs::create_dir_all(&proto_out_dir).unwrap();

    protobuf_codegen_pure::run(protobuf_codegen_pure::Args {
        out_dir: &proto_out_dir.to_string_lossy(),
        includes: &[
            "vendor/grr/grr/proto",
            "vendor/protobuf/src"
        ],
        input: PROTOS,
        ..Default::default()
    }).unwrap();

    std::fs::write(proto_out_dir.join("mod.rs"), b"
        pub mod jobs;
        pub mod sysinfo;
        pub mod timeline;
        pub mod user;
        pub mod knowledge_base;
    ").unwrap();
}
<commit_msg>Add more generated modules to the build script.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::path::PathBuf;

const PROTOS: &'static [&'static str] = &[
    "vendor/grr/grr/proto/grr_response_proto/semantic.proto",
    "vendor/grr/grr/proto/grr_response_proto/sysinfo.proto",
    "vendor/grr/grr/proto/grr_response_proto/knowledge_base.proto",
    "vendor/grr/grr/proto/grr_response_proto/jobs.proto",
    "vendor/grr/grr/proto/grr_response_proto/timeline.proto",
    "vendor/grr/grr/proto/grr_response_proto/anomaly.proto",
    "vendor/grr/grr/proto/grr_response_proto/export.proto",
    "vendor/grr/grr/proto/grr_response_proto/objects.proto",
    "vendor/grr/grr/proto/grr_response_proto/output_plugin.proto",
    "vendor/grr/grr/proto/grr_response_proto/flows.proto",
    "vendor/grr/grr/proto/grr_response_proto/sysinfo.proto",
    "vendor/grr/grr/proto/grr_response_proto/user.proto",
];

const INCLUDES: &'static [&'static str] = &[
    "vendor/grr/grr/proto",
];

fn main() {
    // `prost` generated code.
    prost_build::compile_protos(&PROTOS, &INCLUDES)
        .expect("failed to compile proto files");

    // There is a problem with one enum generated by PROST!: it's values use
    // name mangling, but it's default value does not. This is likely a bug in
    // PROST! itself, but for now we hack around it by replacing the spurious
    // line in the output file ourselves.
    let outdir: PathBuf = std::env::var("OUT_DIR")
        .expect("no output directory")
        .into();

    let target = outdir.join("grr.rs");

    let grr = std::fs::read_to_string(&target)
        .expect("failed to read the generated output file")
        .replace("TskFsAttrTypeDefault", "Default");

    std::fs::write(&target, grr)
        .expect("failed to write updated output file");

    // `protobuf` generated code.
    let proto_out_dir = outdir.join("proto");
    std::fs::create_dir_all(&proto_out_dir).unwrap();

    protobuf_codegen_pure::run(protobuf_codegen_pure::Args {
        out_dir: &proto_out_dir.to_string_lossy(),
        includes: &[
            "vendor/grr/grr/proto",
            "vendor/protobuf/src"
        ],
        input: PROTOS,
        ..Default::default()
    }).unwrap();

    std::fs::write(proto_out_dir.join("mod.rs"), b"
        pub mod anomaly;
        pub mod jobs;
        pub mod export;
        pub mod flows;
        pub mod objects;
        pub mod output_plugin;
        pub mod sysinfo;
        pub mod timeline;
        pub mod user;
        pub mod knowledge_base;
    ").unwrap();
}
<|endoftext|>