<commit_before>pub use self::os::{FILE};
pub use self::os::{fpos_t};
pub use self::os::{BUFSIZ};
pub use self::os::{L_ctermid};
pub use self::os::{L_tmpnam};
pub use self::os::{_IOFBF};
pub use self::os::{_IOLBF};
pub use self::os::{_IONBF};
pub use self::os::{SEEK_CUR};
pub use self::os::{SEEK_END};
pub use self::os::{SEEK_SET};
pub use self::os::{FILENAME_MAX};
pub use self::os::{FOPEN_MAX};
pub use self::os::{TMP_MAX};
pub use self::os::{EOF};
pub use self::os::{stdin};
pub use self::os::{stdout};
pub use self::os::{stderr};

#[cfg(target_os = "linux")]
#[path = "linux/mod.rs"]
mod os;

extern "C" {
    pub fn clearerr(stream: *mut FILE);
    pub fn ctermid(s: *mut ::char_t) -> *mut ::char_t;
    pub fn dprintf(fd: ::int_t, fmt: *const ::char_t, ...) -> ::int_t;
    pub fn fclose(stream: *mut FILE) -> ::int_t;
    pub fn fdopen(fd: ::int_t, modes: *const ::char_t) -> *mut FILE;
    pub fn feof(stream: *mut FILE) -> ::int_t;
    pub fn ferror(stream: *mut FILE) -> ::int_t;
    pub fn fflush(stream: *mut FILE) -> ::int_t;
    pub fn fgetc(stream: *mut FILE) -> ::int_t;
    pub fn fgetpos(stream: *mut FILE, pos: *mut fpos_t) -> ::int_t;
    pub fn fgets(s: *mut ::char_t, n: ::int_t, stream: *mut FILE) -> *mut ::char_t;
    pub fn fileno(stream: *mut FILE) -> ::int_t;
    pub fn flockfile(stream: *mut FILE);
    pub fn fmemopen(s: *mut ::void_t, len: ::size_t, modes: *const ::char_t) -> *mut FILE;
    pub fn fopen(filename: *const ::char_t, modes: *const ::char_t) -> *mut FILE;
    pub fn fprintf(stream: *mut FILE, format: *const ::char_t, ...) -> ::int_t;
    pub fn fputc(c: ::int_t, stream: *mut FILE) -> ::int_t;
    pub fn fputs(s: *const ::char_t, stream: *mut FILE) -> ::int_t;
    pub fn fread(ptr: *mut ::void_t, size: ::size_t, n: ::size_t, stream: *mut FILE) -> ::size_t;
    pub fn fread_unlocked(ptr: *mut ::void_t, size: ::size_t, n: ::size_t, stream: *mut FILE) -> ::size_t;
    pub fn freopen(filename: *const ::char_t, modes: *const ::char_t, stream: *mut FILE) -> *mut FILE;
    pub fn fscanf(stream: *mut FILE, format: *const ::char_t, ...) -> ::int_t;
    pub fn fseeko(stream: *mut FILE, off: ::sys::types::off_t, whence: ::int_t) -> ::int_t;
    pub fn fseek(stream: *mut FILE, off: ::long_t, whence: ::int_t) -> ::int_t;
    pub fn fsetpos(stream: *mut FILE, pos: *const fpos_t) -> ::int_t;
    pub fn ftello(stream: *mut FILE) -> ::sys::types::off_t;
    pub fn ftell(stream: *mut FILE) -> ::long_t;
    pub fn ftrylockfile(stream: *mut FILE) -> ::int_t;
    pub fn funlockfile(stream: *mut FILE);
    pub fn fwrite(ptr: *const ::void_t, size: ::size_t, n: ::size_t, s: *mut FILE) -> ::size_t;
    pub fn getchar() -> ::int_t;
    pub fn getchar_unlocked() -> ::int_t;
    pub fn getc(stream: *mut FILE) -> ::int_t;
    pub fn getc_unlocked(stream: *mut FILE) -> ::int_t;
    pub fn getdelim(lineptr: *mut *mut ::char_t, n: *mut ::size_t, delimiter: ::int_t, stream: *mut FILE) -> ::ssize_t;
    pub fn getline(lineptr: *mut *mut ::char_t, n: *mut ::size_t, stream: *mut FILE) -> ::ssize_t;
    pub fn gets(s: *mut ::char_t) -> *mut ::char_t;
    pub fn getw(stream: *mut FILE) -> ::int_t;
    pub fn open_memstream(bufloc: *mut *mut ::char_t, sizeloc: *mut ::size_t) -> *mut FILE;
    pub fn pclose(stream: *mut FILE) -> ::int_t;
    pub fn perror(s: *const ::char_t);
    pub fn popen(command: *const ::char_t, modes: *const ::char_t) -> *mut FILE;
    pub fn printf(format: *const ::char_t, ...) -> ::int_t;
    pub fn putc(c: ::int_t, stream: *mut FILE) -> ::int_t;
    pub fn putchar(c: ::int_t) -> ::int_t;
    pub fn putchar_unlocked(c: ::int_t) -> ::int_t;
    pub fn putc_unlocked(c: ::int_t, stream: *mut FILE) -> ::int_t;
    pub fn puts(s: *const ::char_t) -> ::int_t;
    pub fn putw(w: ::int_t, stream: *mut FILE) -> ::int_t;
    pub fn remove(filename: *const ::char_t) -> ::int_t;
    pub fn renameat(oldfd: ::int_t, old: *const ::char_t, newfd: ::int_t, new: *const ::char_t) -> ::int_t;
    pub fn rename(old: *const ::char_t, new: *const ::char_t) -> ::int_t;
    pub fn rewind(stream: *mut FILE);
    pub fn scanf(format: *const ::char_t, ...) -> ::int_t;
    pub fn setbuf(stream: *mut FILE, buf: *mut ::char_t);
    pub fn setvbuf(stream: *mut FILE, buf: *mut ::char_t, modes: ::int_t, n: ::size_t) -> ::int_t;
    pub fn snprintf(s: *mut ::char_t, maxlen: ::size_t, format: *const ::char_t, ...) -> ::int_t;
    pub fn sprintf(s: *mut ::char_t, format: *const ::char_t, ...) -> ::int_t;
    pub fn sscanf(s: *const ::char_t, format: *const ::char_t, ...) -> ::int_t;
    pub fn tempnam(dir: *const ::char_t, pfx: *const ::char_t) -> *mut ::char_t;
    pub fn tmpfile() -> *mut FILE;
    pub fn tmpnam(s: *mut ::char_t) -> *mut ::char_t;
    pub fn ungetc(c: ::int_t, stream: *mut FILE) -> ::int_t;
}
<commit_msg>stdio.h<commit_after>pub use self::os::{FILE};
pub use self::os::{fpos_t};
pub use self::os::{BUFSIZ};
pub use self::os::{L_ctermid};
pub use self::os::{L_tmpnam};
pub use self::os::{_IOFBF};
pub use self::os::{_IOLBF};
pub use self::os::{_IONBF};
pub use self::os::{SEEK_CUR};
pub use self::os::{SEEK_END};
pub use self::os::{SEEK_SET};
pub use self::os::{FILENAME_MAX};
pub use self::os::{FOPEN_MAX};
pub use self::os::{TMP_MAX};
pub use self::os::{EOF};
pub use self::os::{stdin};
pub use self::os::{stdout};
pub use self::os::{stderr};

use {NTStr, char_t, size_t, int_t, ssize_t, void_t, long_t};
use sys::types::{off_t};

#[cfg(target_os = "linux")]
#[path = "linux/mod.rs"]
mod os;

pub fn clearerr(stream: *mut FILE) {
    extern { fn clearerr(stream: *mut FILE); }
    unsafe { clearerr(stream); }
}

extern {
    pub fn ctermid(s: *mut char_t) -> *mut char_t;
}

extern {
    pub fn dprintf(fd: int_t, fmt: *const char_t, ...) -> int_t;
}

pub fn fclose(stream: *mut FILE) -> int_t {
    extern { fn fclose(stream: *mut FILE) -> int_t; }
    unsafe { fclose(stream) }
}

pub fn fdopen<T: NTStr>(fd: int_t, modes: &T) -> *mut FILE {
    extern { fn fdopen(fd: int_t, modes: *const char_t) -> *mut FILE; }
    unsafe { fdopen(fd, modes.as_ptr()) }
}

pub fn feof(stream: *mut FILE) -> int_t {
    extern { fn feof(stream: *mut FILE) -> int_t; }
    unsafe { feof(stream) }
}

pub fn ferror(stream: *mut FILE) -> int_t {
    extern { fn ferror(stream: *mut FILE) -> int_t; }
    unsafe { ferror(stream) }
}

pub fn fflush(stream: *mut FILE) -> int_t {
    extern { fn fflush(stream: *mut FILE) -> int_t; }
    unsafe { fflush(stream) }
}

pub fn fgetc(stream: *mut FILE) -> int_t {
    extern { fn fgetc(stream: *mut FILE) -> int_t; }
    unsafe { fgetc(stream) }
}

pub fn fgetpos(stream: *mut FILE, pos: &mut fpos_t) -> int_t {
    extern { fn fgetpos(stream: *mut FILE, pos: *mut fpos_t) -> int_t; }
    unsafe { fgetpos(stream, pos as *mut _) }
}

pub fn fgets(s: &mut [u8], stream: *mut FILE) -> *mut char_t {
    extern { fn fgets(s: *mut char_t, n: int_t, stream: *mut FILE) -> *mut char_t; }
    unsafe { fgets(s.as_mut_ptr() as *mut _, s.len() as int_t, stream) }
}

pub fn fileno(stream: *mut FILE) -> int_t {
    extern { fn fileno(stream: *mut FILE) -> int_t; }
    unsafe { fileno(stream) }
}

pub fn flockfile(stream: *mut FILE) {
    extern { fn flockfile(stream: *mut FILE); }
    unsafe { flockfile(stream); }
}

pub fn fmemopen(s: &mut [u8], modes: *const char_t) -> *mut FILE {
    extern { fn fmemopen(s: *mut void_t, len: size_t, modes: *const char_t) -> *mut FILE; }
    unsafe { fmemopen(s.as_ptr() as *mut _, s.len() as size_t, modes) }
}

pub fn fopen<T: NTStr, U: NTStr>(filename: &T, modes: &U) -> *mut FILE {
    extern { fn fopen(filename: *const char_t, modes: *const char_t) -> *mut FILE; }
    unsafe { fopen(filename.as_ptr(), modes.as_ptr()) }
}

extern {
    pub fn fprintf(stream: *mut FILE, format: *const char_t, ...) -> int_t;
}

pub fn fputc(c: int_t, stream: *mut FILE) -> int_t {
    extern { fn fputc(c: int_t, stream: *mut FILE) -> int_t; }
    unsafe { fputc(c, stream) }
}

pub fn fputs<T: NTStr>(s: &T, stream: *mut FILE) -> int_t {
    extern { fn fputs(s: *const char_t, stream: *mut FILE) -> int_t; }
    unsafe { fputs(s.as_ptr(), stream) }
}

extern {
    pub fn fread(ptr: *mut void_t, size: size_t, n: size_t, stream: *mut FILE) -> size_t;
    pub fn fread_unlocked(ptr: *mut void_t, size: size_t, n: size_t,
                          stream: *mut FILE) -> size_t;
}

pub fn freopen<T: NTStr, U: NTStr>(filename: &T, modes: &U,
                                   stream: *mut FILE) -> *mut FILE {
    extern { fn freopen(filename: *const char_t, modes: *const char_t,
                        stream: *mut FILE) -> *mut FILE; }
    unsafe { freopen(filename.as_ptr(), modes.as_ptr(), stream) }
}

extern {
    pub fn fscanf(stream: *mut FILE, format: *const char_t, ...) -> int_t;
}

pub fn fseeko(stream: *mut FILE, off: off_t, whence: int_t) -> int_t {
    extern { fn fseeko(stream: *mut FILE, off: off_t, whence: int_t) -> int_t; }
    unsafe { fseeko(stream, off, whence) }
}

pub fn fseek(stream: *mut FILE, off: long_t, whence: int_t) -> int_t {
    extern { fn fseek(stream: *mut FILE, off: long_t, whence: int_t) -> int_t; }
    unsafe { fseek(stream, off, whence) }
}

pub fn fsetpos(stream: *mut FILE, pos: &fpos_t) -> int_t {
    extern { fn fsetpos(stream: *mut FILE, pos: *const fpos_t) -> int_t; }
    unsafe { fsetpos(stream, pos as *const _) }
}

pub fn ftello(stream: *mut FILE) -> off_t {
    extern { fn ftello(stream: *mut FILE) -> off_t; }
    unsafe { ftello(stream) }
}

pub fn ftell(stream: *mut FILE) -> long_t {
    extern { fn ftell(stream: *mut FILE) -> long_t; }
    unsafe { ftell(stream) }
}

pub fn ftrylockfile(stream: *mut FILE) -> int_t {
    extern { fn ftrylockfile(stream: *mut FILE) -> int_t; }
    unsafe { ftrylockfile(stream) }
}

pub fn funlockfile(stream: *mut FILE) {
    extern { fn funlockfile(stream: *mut FILE); }
    unsafe { funlockfile(stream); }
}

extern {
    pub fn fwrite(ptr: *const void_t, size: size_t, n: size_t, s: *mut FILE) -> size_t;
}

pub fn getchar() -> int_t {
    extern { fn getchar() -> int_t; }
    unsafe { getchar() }
}

pub fn getchar_unlocked() -> int_t {
    extern { fn getchar_unlocked() -> int_t; }
    unsafe { getchar_unlocked() }
}

pub fn getc(stream: *mut FILE) -> int_t {
    extern { fn getc(stream: *mut FILE) -> int_t; }
    unsafe { getc(stream) }
}

pub fn getc_unlocked(stream: *mut FILE) -> int_t {
    extern { fn getc_unlocked(stream: *mut FILE) -> int_t; }
    unsafe { getc_unlocked(stream) }
}

extern {
    pub fn getdelim(lineptr: *mut *mut char_t, n: *mut size_t, delimiter: int_t,
                    stream: *mut FILE) -> ssize_t;
    pub fn getline(lineptr: *mut *mut char_t, n: *mut size_t, stream: *mut FILE) -> ssize_t;
    pub fn gets(s: *mut char_t) -> *mut char_t;
}

pub fn getw(stream: *mut FILE) -> int_t {
    extern { fn getw(stream: *mut FILE) -> int_t; }
    unsafe { getw(stream) }
}

extern {
    pub fn open_memstream(bufloc: *mut *mut char_t, sizeloc: *mut size_t) -> *mut FILE;
}

pub fn pclose(stream: *mut FILE) -> int_t {
    extern { fn pclose(stream: *mut FILE) -> int_t; }
    unsafe { pclose(stream) }
}

pub fn perror<T: NTStr>(s: &T) {
    extern { fn perror(s: *const char_t); }
    unsafe { perror(s.as_ptr()); }
}

pub fn popen<T: NTStr, U: NTStr>(command: &T, modes: &U) -> *mut FILE {
    extern { fn popen(command: *const char_t, modes: *const char_t) -> *mut FILE; }
    unsafe { popen(command.as_ptr(), modes.as_ptr()) }
}

extern {
    pub fn printf(format: *const char_t, ...) -> int_t;
}

pub fn putc(c: int_t, stream: *mut FILE) -> int_t {
    extern { fn putc(c: int_t, stream: *mut FILE) -> int_t; }
    unsafe { putc(c, stream) }
}

pub fn putchar(c: int_t) -> int_t {
    extern { fn putchar(c: int_t) -> int_t; }
    unsafe { putchar(c) }
}

pub fn putchar_unlocked(c: int_t) -> int_t {
    extern { fn putchar_unlocked(c: int_t) -> int_t; }
    unsafe { putchar_unlocked(c) }
}

pub fn putc_unlocked(c: int_t, stream: *mut FILE) -> int_t {
    extern { fn putc_unlocked(c: int_t, stream: *mut FILE) -> int_t; }
    unsafe { putc_unlocked(c, stream) }
}

pub fn puts<T: NTStr>(s: &T) -> int_t {
    extern { fn puts(s: *const char_t) -> int_t; }
    unsafe { puts(s.as_ptr()) }
}

pub fn putw(w: int_t, stream: *mut FILE) -> int_t {
    extern { fn putw(w: int_t, stream: *mut FILE) -> int_t; }
    unsafe { putw(w, stream) }
}

pub fn remove<T: NTStr>(filename: &T) -> int_t {
    extern { fn remove(filename: *const char_t) -> int_t; }
    unsafe { remove(filename.as_ptr()) }
}

pub fn renameat<T: NTStr, U: NTStr>(oldfd: int_t, old: &T, newfd: int_t,
                                    new: &T) -> int_t {
    extern { fn renameat(oldfd: int_t, old: *const char_t, newfd: int_t,
                         new: *const char_t) -> int_t; }
    unsafe { renameat(oldfd, old.as_ptr(), newfd, new.as_ptr()) }
}

pub fn rename<T: NTStr, U: NTStr>(old: &T, new: &U) -> int_t {
    extern { fn rename(old: *const char_t, new: *const char_t) -> int_t; }
    unsafe { rename(old.as_ptr(), new.as_ptr()) }
}

pub fn rewind(stream: *mut FILE) {
    extern { fn rewind(stream: *mut FILE); }
    unsafe { rewind(stream); }
}

extern {
    pub fn scanf(format: *const char_t, ...) -> int_t;
}

extern {
    pub fn setbuf(stream: *mut FILE, buf: *mut char_t);
    pub fn setvbuf(stream: *mut FILE, buf: *mut char_t, modes: int_t, n: size_t) -> int_t;
    pub fn snprintf(s: *mut char_t, maxlen: size_t, format: *const char_t, ...) -> int_t;
    pub fn sprintf(s: *mut char_t, format: *const char_t, ...) -> int_t;
    pub fn sscanf(s: *const char_t, format: *const char_t, ...) -> int_t;
    pub fn tempnam(dir: *const char_t, pfx: *const char_t) -> *mut char_t;
}

pub fn tmpfile() -> *mut FILE {
    extern { fn tmpfile() -> *mut FILE; }
    unsafe { tmpfile() }
}

extern {
    pub fn tmpnam(s: *mut char_t) -> *mut char_t;
}

pub fn ungetc(c: int_t, stream: *mut FILE) -> int_t {
    extern { fn ungetc(c: int_t, stream: *mut FILE) -> int_t; }
    unsafe { ungetc(c, stream) }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol,
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::{HandshakeResult, MidHandshakeTlsStream, TcpStream},
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, info, trace};
use std::{
    collections::VecDeque,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    SocketConnected,
    SocketWritable,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    handshake: Option<MidHandshakeTlsStream>,
    stream: Option<TcpStream>,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let (mut stream, mut handshake) = match stream {
            Ok(stream) => (Some(stream), None),
            Err(error) => (None, Some(error.into_mid_handshake_tls_stream()?)),
        };
        let source = stream
            .as_mut()
            .unwrap_or_else(|| handshake.as_mut().unwrap().get_mut());
        let slot = reactor.register(source, socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            handshake,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn ensure_connected(&mut self) -> Result<bool> {
        match self
            .stream
            .as_ref()
            .unwrap_or_else(|| self.handshake.as_ref().unwrap().get_ref())
            .peer_addr()
        {
            Ok(peer) => {
                info!("Connected to {}", peer);
                self.status = Status::SocketConnected;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected = err.kind() {
                    Ok(false)
                } else {
                    Err(err)?
                }
            }
        }
    }

    fn ensure_writable(&mut self) -> Result<bool> {
        match self
            .stream
            .as_ref()
            .unwrap_or_else(|| self.handshake.as_ref().unwrap().get_ref())
            .is_writable()
        {
            Ok(()) => {
                self.status = Status::SocketWritable;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected | io::ErrorKind::WouldBlock = err.kind() {
                    Ok(false)
                } else {
                    Err(err)?
                }
            }
        }
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        loop {
            let success = match self.status {
                Status::Initial => self.ensure_connected(),
                Status::SocketConnected => self.ensure_writable(),
                Status::SocketWritable => return self.finish_setup(),
                Status::Connected => return Ok(true),
                Status::Stop => return Ok(false),
            }?;
            if !success {
                return Ok(false);
            }
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable()
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown();
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if let Some(handshake) = self.handshake.take() {
            match handshake.handshake() {
                Ok(stream) => self.stream = Some(stream),
                Err(error) => self.handshake = Some(error.into_mid_handshake_tls_stream()?),
            }
            if self.handshake.is_some() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.reactor.shutdown();
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        self.stream.as_mut().unwrap().flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        if self.socket_state.writable() {
            let res = self.flush();
            self.handle_write_result(res)?;
        }
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            if let Err(e) = self
                .socket_state
                .handle_read_result(res, &*self.reactor, self.slot)
            {
                error!("error reading: {:?}", e);
                self.critical_error(e)?;
            }
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(self.stream.as_mut().unwrap())?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(self.stream.as_mut().unwrap())
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz);
                    })?;
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<commit_msg>refuse frames too large<commit_after>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol::{self, AMQPError, AMQPHardError},
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::{HandshakeResult, MidHandshakeTlsStream, TcpStream},
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, info, trace};
use std::{
    collections::VecDeque,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    SocketConnected,
    SocketWritable,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    handshake: Option<MidHandshakeTlsStream>,
    stream: Option<TcpStream>,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let (mut stream, mut handshake) = match stream {
            Ok(stream) => (Some(stream), None),
            Err(error) => (None, Some(error.into_mid_handshake_tls_stream()?)),
        };
        let source = stream
            .as_mut()
            .unwrap_or_else(|| handshake.as_mut().unwrap().get_mut());
        let slot = reactor.register(source, socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            handshake,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn ensure_connected(&mut self) -> Result<bool> {
        match self
            .stream
            .as_ref()
            .unwrap_or_else(|| self.handshake.as_ref().unwrap().get_ref())
            .peer_addr()
        {
            Ok(peer) => {
                info!("Connected to {}", peer);
                self.status = Status::SocketConnected;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected = err.kind() {
                    Ok(false)
                } else {
                    Err(err)?
                }
            }
        }
    }

    fn ensure_writable(&mut self) -> Result<bool> {
        match self
            .stream
            .as_ref()
            .unwrap_or_else(|| self.handshake.as_ref().unwrap().get_ref())
            .is_writable()
        {
            Ok(()) => {
                self.status = Status::SocketWritable;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected | io::ErrorKind::WouldBlock = err.kind() {
                    Ok(false)
                } else {
                    Err(err)?
                }
            }
        }
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        loop {
            let success = match self.status {
                Status::Initial => self.ensure_connected(),
                Status::SocketConnected => self.ensure_writable(),
                Status::SocketWritable => return self.finish_setup(),
                Status::Connected => return Ok(true),
                Status::Stop => return Ok(false),
            }?;
            if !success {
                return Ok(false);
            }
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable()
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown();
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if let Some(handshake) = self.handshake.take() {
            match handshake.handshake() {
                Ok(stream) => self.stream = Some(stream),
                Err(error) => self.handshake = Some(error.into_mid_handshake_tls_stream()?),
            }
            if self.handshake.is_some() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.reactor.shutdown();
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        self.stream.as_mut().unwrap().flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        if self.socket_state.writable() {
            let res = self.flush();
            self.handle_write_result(res)?;
        }
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            if let Err(e) = self
                .socket_state
                .handle_read_result(res, &*self.reactor, self.slot)
            {
                error!("error reading: {:?}", e);
                self.critical_error(e)?;
            }
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(self.stream.as_mut().unwrap())?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(self.stream.as_mut().unwrap())
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz);
                    })?;
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                let frame_max = self.configuration.frame_max() as usize;
                if frame_max > 0 && consumed > frame_max {
                    error!("received large ({} bytes) frame", consumed);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.handle().close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.poll_internal_rpc()?;
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>


<commit_before>#![crate_name = "tail"]
/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Morten Olsen Lysgaard <morten@lysgaard.no>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 *
 */

#![feature(macro_rules)]

extern crate getopts;

use std::char::UnicodeChar;
use std::io::{stdin};
use std::io::{BufferedReader, BytesReader};
use std::io::fs::File;
use std::path::Path;
use std::str::from_utf8;
use getopts::{optopt, optflag, getopts, usage};
use std::collections::ring_buf::RingBuf;
use std::io::timer::sleep;
use std::time::duration::Duration;

#[path = "../common/util.rs"]
mod util;

static NAME: &'static str = "tail";
static VERSION: &'static str = "0.0.1";

pub fn uumain(args: Vec<String>) -> int {
    let mut beginning = false;
    let mut lines = true;
    let mut byte_count = 0u;
    let mut line_count = 10u;
    let mut sleep_msec = 1000u64;

    // handle obsolete -number syntax
    let options = match obsolete(args.tail()) {
        (args, Some(n)) => { line_count = n; args },
        (args, None) => args
    };

    let args = options;

    let possible_options = [
        optopt("c", "bytes", "Number of bytes to print", "k"),
        optopt("n", "lines", "Number of lines to print", "k"),
        optflag("f", "follow", "Print the file as it grows"),
        optopt("s", "sleep-interval", "Number or seconds to sleep between polling the file when running with -f", "n"),
        optflag("h", "help", "help"),
        optflag("V", "version", "version"),
    ];

    let given_options = match getopts(args.as_slice(), &possible_options) {
        Ok (m) => { m }
        Err(_) => {
            println!("{}", usage(NAME, &possible_options));
            return 1;
        }
    };

    if given_options.opt_present("h") {
        println!("{}", usage(NAME, &possible_options));
        return 0;
    }
    if given_options.opt_present("V") { version(); return 0 }

    let follow = given_options.opt_present("f");
    if follow {
        match given_options.opt_str("s") {
            Some(n) => {
                let parsed: Option<u64> = from_str(n.as_slice());
                match parsed {
                    Some(m) => { sleep_msec = m * 1000 }
                    None => {}
                }
            }
            None => {}
        };
    }

    match given_options.opt_str("n") {
        Some(n) => {
            let mut slice = n.as_slice();
            if slice.len() > 0 && slice.char_at(0) == '+' {
                beginning = true;
                slice = slice.slice_from(1);
            }
            line_count = match from_str(slice) {
                Some(m) => m,
                None => {
                    show_error!("invalid number of lines ({})", slice);
                    return 1;
                }
            };
        }
        None => match given_options.opt_str("c") {
            Some(n) => {
                let mut slice = n.as_slice();
                if slice.len() > 0 && slice.char_at(0) == '+' {
                    beginning = true;
                    slice = slice.slice_from(1);
                }
                byte_count = match from_str(slice) {
                    Some(m) => m,
                    None => {
                        show_error!("invalid number of bytes ({})", slice);
                        return 1;
                    }
                };
                lines = false;
            }
            None => { }
        }
    };

    let files = given_options.free;

    if files.is_empty() {
        let mut buffer = BufferedReader::new(stdin());
        tail(&mut buffer, line_count, byte_count, beginning, lines, follow, sleep_msec);
    } else {
        let mut multiple = false;
        let mut firstime = true;

        if files.len() > 1 {
            multiple = true;
        }

        for file in files.iter() {
            if multiple {
                if !firstime { println!(""); }
                println!("==> {} <==", file.as_slice());
            }
            firstime = false;

            let path = Path::new(file.as_slice());
            let reader = File::open(&path).unwrap();
            let mut buffer = BufferedReader::new(reader);
            tail(&mut buffer, line_count, byte_count, beginning, lines, follow, sleep_msec);
        }
    }

    0
}

// It searches for an option in the form of -123123
//
// In case is found, the options vector will get rid of that object so that
// getopts works correctly.
fn obsolete(options: &[String]) -> (Vec<String>, Option<uint>) {
    let mut options: Vec<String> = options.to_vec();
    let mut a = 0;
    let b = options.len();

    while a < b {
        let current = options[a].clone();
        let current = current.as_bytes();

        if current.len() > 1 && current[0] == '-' as u8 {
            let len = current.len();
            for pos in range(1, len) {
                // Ensure that the argument is only made out of digits
                if !UnicodeChar::is_numeric(current[pos] as char) { break; }

                // If this is the last number
                if pos == len - 1 {
                    options.remove(a);
                    let number: Option<uint> = from_str(from_utf8(current.slice(1,len)).unwrap());
                    return (options, Some(number.unwrap()));
                }
            }
        }

        a += 1;
    };

    (options, None)
}

fn tail<T: Reader>(reader: &mut BufferedReader<T>, line_count: uint, byte_count: uint, beginning: bool, lines: bool, follow: bool, sleep_msec: u64) {
    if lines {
        tail_lines(reader, line_count, beginning);
    } else {
        tail_bytes(reader, byte_count, beginning);
    }

    // if we follow the file, sleep a bit and print the rest if the file has grown.
    while follow {
        sleep(Duration::milliseconds(sleep_msec as i64));
        for io_line in reader.lines() {
            match io_line {
                Ok(line) => print!("{}", line),
                Err(err) => panic!(err)
            }
        }
    }
}

#[inline]
fn tail_lines<T: Reader>(reader: &mut BufferedReader<T>, mut line_count: uint, beginning: bool) {
    // read through each line and store them in a ringbuffer that always contains
    // line_count lines. When reaching the end of file, output the lines in the
    // ringbuf.
    let mut ringbuf: RingBuf<String> = RingBuf::new();
    let mut lines = reader.lines().skip(
        if beginning {
            let temp = line_count;
            line_count = ::std::uint::MAX;
            temp - 1
        } else {
            0
        }
    );
    for io_line in lines {
        match io_line {
            Ok(line) => {
                if line_count <= ringbuf.len() {
                    ringbuf.pop_front();
                }
                ringbuf.push_back(line);
            }
            Err(err) => panic!(err)
        }
    }
    for line in ringbuf.iter() {
        print!("{}", line);
    }
}

#[inline]
fn tail_bytes<T: Reader>(reader: &mut BufferedReader<T>, mut byte_count: uint, beginning: bool) {
    let mut ringbuf: RingBuf<u8> = RingBuf::new();
    let mut bytes = reader.bytes().skip(
        if beginning {
            let temp = byte_count;
            byte_count = ::std::uint::MAX;
            temp - 1
        } else {
            0
        }
    );
    for io_byte in bytes {
        match io_byte {
            Ok(byte) => {
                if byte_count <= ringbuf.len() {
                    ringbuf.pop_front();
                }
                ringbuf.push_back(byte);
            }
            Err(err) => panic!(err)
        }
    }
    for byte in ringbuf.iter() {
        print!("{}", byte);
    }
}

fn version () {
    println!("{} v{}", NAME, VERSION);
}
<commit_msg>tail: print bytes as char instead of u8<commit_after>#![crate_name = "tail"]
/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Morten Olsen Lysgaard <morten@lysgaard.no>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 *
 */

#![feature(macro_rules)]

extern crate getopts;

use std::char::UnicodeChar;
use std::io::{stdin};
use std::io::{BufferedReader, BytesReader};
use std::io::fs::File;
use std::path::Path;
use std::str::from_utf8;
use getopts::{optopt, optflag, getopts, usage};
use std::collections::ring_buf::RingBuf;
use std::io::timer::sleep;
use std::time::duration::Duration;

#[path = "../common/util.rs"]
mod util;

static NAME: &'static str = "tail";
static VERSION: &'static str = "0.0.1";

pub fn uumain(args: Vec<String>) -> int {
    let mut beginning = false;
    let mut lines = true;
    let mut byte_count = 0u;
    let mut line_count = 10u;
    let mut sleep_msec = 1000u64;

    // handle obsolete -number syntax
    let options = match obsolete(args.tail()) {
        (args, Some(n)) => { line_count = n; args },
        (args, None) => args
    };

    let args = options;

    let possible_options = [
        optopt("c", "bytes", "Number of bytes to print", "k"),
        optopt("n", "lines", "Number of lines to print", "k"),
        optflag("f", "follow", "Print the file as it grows"),
        optopt("s", "sleep-interval", "Number or seconds to sleep between polling the file when running with -f", "n"),
        optflag("h", "help", "help"),
        optflag("V", "version", "version"),
    ];

    let given_options = match getopts(args.as_slice(), &possible_options) {
        Ok (m) => { m }
        Err(_) => {
            println!("{}", usage(NAME, &possible_options));
            return 1;
        }
    };

    if given_options.opt_present("h") {
        println!("{}", usage(NAME, &possible_options));
        return 0;
    }
    if given_options.opt_present("V") { version(); return 0 }

    let follow = given_options.opt_present("f");
    if follow {
        match given_options.opt_str("s") {
            Some(n) => {
                let parsed: Option<u64> = from_str(n.as_slice());
                match parsed {
                    Some(m) => { sleep_msec = m * 1000 }
                    None => {}
                }
            }
            None => {}
        };
    }

    match given_options.opt_str("n") {
        Some(n) => {
            let mut slice = n.as_slice();
            if slice.len() > 0 && slice.char_at(0) == '+' {
                beginning = true;
                slice = slice.slice_from(1);
            }
            line_count = match from_str(slice) {
                Some(m) => m,
                None => {
                    show_error!("invalid number of lines ({})", slice);
                    return 1;
                }
            };
        }
        None => match given_options.opt_str("c") {
            Some(n) => {
                let mut slice = n.as_slice();
                if slice.len() > 0 && slice.char_at(0) == '+' {
                    beginning = true;
                    slice = slice.slice_from(1);
                }
                byte_count = match from_str(slice) {
                    Some(m) => m,
                    None => {
                        show_error!("invalid number of bytes ({})", slice);
                        return 1;
                    }
                };
                lines = false;
            }
            None => { }
        }
    };

    let files = given_options.free;

    if files.is_empty() {
        let mut buffer = BufferedReader::new(stdin());
        tail(&mut buffer, line_count, byte_count, beginning, lines, follow, sleep_msec);
    } else {
        let mut multiple = false;
        let mut firstime = true;

        if files.len() > 1 {
            multiple = true;
        }

        for file in files.iter() {
            if multiple {
                if !firstime { println!(""); }
                println!("==> {} <==", file.as_slice());
            }
            firstime = false;

            let path = Path::new(file.as_slice());
            let reader = File::open(&path).unwrap();
            let mut buffer = BufferedReader::new(reader);
            tail(&mut buffer, line_count, byte_count, beginning, lines, follow, sleep_msec);
        }
    }

    0
}

// It searches for an option in the form of -123123
//
// In case is found, the options vector will get rid of that object so that
// getopts works correctly.
fn obsolete(options: &[String]) -> (Vec<String>, Option<uint>) {
    let mut options: Vec<String> = options.to_vec();
    let mut a = 0;
    let b = options.len();

    while a < b {
        let current = options[a].clone();
        let current = current.as_bytes();

        if current.len() > 1 && current[0] == '-' as u8 {
            let len = current.len();
            for pos in range(1, len) {
                // Ensure that the argument is only made out of digits
                if !UnicodeChar::is_numeric(current[pos] as char) { break; }

                // If this is the last number
                if pos == len - 1 {
                    options.remove(a);
                    let number: Option<uint> = from_str(from_utf8(current.slice(1,len)).unwrap());
                    return (options, Some(number.unwrap()));
                }
            }
        }

        a += 1;
    };

    (options, None)
}

fn tail<T: Reader>(reader: &mut BufferedReader<T>, line_count: uint, byte_count: uint, beginning: bool, lines: bool, follow: bool, sleep_msec: u64) {
    if lines {
        tail_lines(reader, line_count, beginning);
    } else {
        tail_bytes(reader, byte_count, beginning);
    }

    // if we follow the file, sleep a bit and print the rest if the file has grown.
    while follow {
        sleep(Duration::milliseconds(sleep_msec as i64));
        for io_line in reader.lines() {
            match io_line {
                Ok(line) => print!("{}", line),
                Err(err) => panic!(err)
            }
        }
    }
}

#[inline]
fn tail_lines<T: Reader>(reader: &mut BufferedReader<T>, mut line_count: uint, beginning: bool) {
    // read through each line and store them in a ringbuffer that always contains
    // line_count lines. When reaching the end of file, output the lines in the
    // ringbuf.
    let mut ringbuf: RingBuf<String> = RingBuf::new();
    let mut lines = reader.lines().skip(
        if beginning {
            let temp = line_count;
            line_count = ::std::uint::MAX;
            temp - 1
        } else {
            0
        }
    );
    for io_line in lines {
        match io_line {
            Ok(line) => {
                if line_count <= ringbuf.len() {
                    ringbuf.pop_front();
                }
                ringbuf.push_back(line);
            }
            Err(err) => panic!(err)
        }
    }
    for line in ringbuf.iter() {
        print!("{}", line);
    }
}

#[inline]
fn tail_bytes<T: Reader>(reader: &mut BufferedReader<T>, mut byte_count: uint, beginning: bool) {
    let mut ringbuf: RingBuf<char> = RingBuf::new();
    let mut bytes = reader.bytes().skip(
        if beginning {
            let temp = byte_count;
            byte_count = ::std::uint::MAX;
            temp - 1
        } else {
            0
        }
    );
    for io_byte in bytes {
        match io_byte {
            Ok(byte) => {
                if byte_count <= ringbuf.len() {
                    ringbuf.pop_front();
                }
                ringbuf.push_back(byte as char);
            }
            Err(err) => panic!(err)
        }
    }
    for byte in ringbuf.iter() {
        print!("{}", byte);
    }
}

fn version () {
    println!("{} v{}", NAME, VERSION);
}
<|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/. */

use std::borrow::Cow;
use std::collections::HashMap;
use std::io::Write;

use syn::{self, UnOp};

use bindgen::config::{Config, Language};
use bindgen::declarationtyperesolver::DeclarationTypeResolver;
use bindgen::dependencies::Dependencies;
use bindgen::ir::{
    AnnotationSet, Cfg, ConditionWrite, Documentation, GenericParams, Item, ItemContainer, Path,
    Struct, ToCondition, Type,
};
use bindgen::library::Library;
use bindgen::writer::{Source, SourceWriter};
use bindgen::Bindings;

#[derive(Debug, Clone)]
pub enum Literal {
    Expr(String),
    PostfixUnaryOp {
        op: &'static str,
        value: Box<Literal>,
    },
    BinOp {
        left: Box<Literal>,
        op: &'static str,
        right: Box<Literal>,
    },
    Struct {
        path: Path,
        export_name: String,
        fields: HashMap<String, Literal>,
    },
}

impl Literal {
    fn replace_self_with(&mut self, self_ty: &Path) {
        match *self {
            Literal::PostfixUnaryOp { .. } | Literal::BinOp { .. } | Literal::Expr(..) => {}
            Literal::Struct {
                ref mut path,
                ref mut export_name,
                ref mut fields,
            } => {
                if path.replace_self_with(self_ty) {
                    *export_name = self_ty.name().to_owned();
                }
                for (ref _name, ref mut expr) in fields {
                    expr.replace_self_with(self_ty);
                }
            }
        }
    }

    fn is_valid(&self, bindings: &Bindings) -> bool {
        match *self {
            Literal::Expr(..) => true,
            Literal::PostfixUnaryOp { ref value, .. } => value.is_valid(bindings),
            Literal::BinOp {
                ref left,
                ref right,
                ..
            } => left.is_valid(bindings) && right.is_valid(bindings),
            Literal::Struct { ref path, .. } => bindings.struct_exists(path),
        }
    }
}

impl Literal {
    pub fn rename_for_config(&mut self, config: &Config) {
        match self {
            Literal::Struct {
                ref mut export_name,
                fields,
                ..
            } => {
                config.export.rename(export_name);
                for (_, lit) in fields {
                    lit.rename_for_config(config);
                }
            }
            Literal::PostfixUnaryOp { ref mut value, .. } => {
                value.rename_for_config(config);
            }
            Literal::BinOp {
                ref mut left,
                ref mut right,
                ..
            } => {
                left.rename_for_config(config);
                right.rename_for_config(config);
            }
            Literal::Expr(_) => {}
        }
    }

    // Translate from full blown `syn::Expr` into a simpler `Literal` type
    pub fn load(expr: &syn::Expr) -> Result<Literal, String> {
        match *expr {
            // Match binary expressions of the form `a * b`
            syn::Expr::Binary(ref bin_expr) => {
                let l = Self::load(&bin_expr.left)?;
                let r = Self::load(&bin_expr.right)?;
                let op = match bin_expr.op {
                    syn::BinOp::Add(..) => "+",
                    syn::BinOp::Sub(..) => "-",
                    syn::BinOp::Mul(..) => "*",
                    syn::BinOp::Div(..) => "/",
                    syn::BinOp::Rem(..) => "%",
                    syn::BinOp::Shl(..) => "<<",
                    syn::BinOp::Shr(..) => ">>",
                    _ => return Err(format!("Unsupported binary op {:?}", bin_expr.op)),
                };
                Ok(Literal::BinOp {
                    left: Box::new(l),
                    op,
                    right: Box::new(r),
                })
            }

            // Match literals like "one", 'a', 32 etc
            syn::Expr::Lit(syn::ExprLit { ref lit, .. }) => {
                match lit {
                    syn::Lit::Str(ref value) => {
                        Ok(Literal::Expr(format!("u8\"{}\"", value.value())))
                    }
                    syn::Lit::Byte(ref value) => Ok(Literal::Expr(format!("{}", value.value()))),
                    syn::Lit::Char(ref value) => Ok(Literal::Expr(match value.value() as u32 {
                        0..=255 => format!("'{}'", value.value().escape_default()),
                        other_code => format!(r"L'\u{:X}'", other_code),
                    })),
                    syn::Lit::Int(ref value) => {
                        Ok(Literal::Expr(value.base10_digits().to_string()))
                    }
                    syn::Lit::Float(ref value) => {
                        Ok(Literal::Expr(value.base10_digits().to_string()))
                    }
                    syn::Lit::Bool(ref value) => Ok(Literal::Expr(format!("{}", value.value))),
                    // TODO: Add support for byte string and Verbatim
                    _ => Err(format!("Unsupported literal expression. {:?}", *lit)),
                }
            }

            syn::Expr::Struct(syn::ExprStruct {
                ref path,
                ref fields,
                ..
            }) => {
                let struct_name = path.segments[0].ident.to_string();
                let mut field_map = HashMap::<String, Literal>::default();
                for field in fields {
                    let ident = match field.member {
                        syn::Member::Named(ref name) => name.to_string(),
                        syn::Member::Unnamed(ref index) => format!("_{}", index.index),
                    };
                    let key = ident.to_string();
                    let value = Literal::load(&field.expr)?;
                    field_map.insert(key, value);
                }
                Ok(Literal::Struct {
                    path: Path::new(struct_name.clone()),
                    export_name: struct_name,
                    fields: field_map,
                })
            }

            syn::Expr::Unary(syn::ExprUnary {
                ref op, ref expr, ..
            }) => match *op {
                UnOp::Neg(_) => {
                    let val = Self::load(expr)?;
                    Ok(Literal::PostfixUnaryOp {
                        op: "-",
                        value: Box::new(val),
                    })
                }
                _ => Err(format!("Unsupported Unary expression. {:?}", *op)),
            },

            _ => Err(format!("Unsupported expression. {:?}", *expr)),
        }
    }

    fn write<F: Write>(&self, config: &Config, out: &mut SourceWriter<F>) {
        match self {
            Literal::Expr(v) => write!(out, "{}", v),
            Literal::PostfixUnaryOp { op, ref value } => {
                write!(out, "{}", op);
                value.write(config, out);
            }
            Literal::BinOp {
                ref left,
                op,
                ref right,
            } => {
                left.write(config, out);
                write!(out, " {} ", op);
                right.write(config, out);
            }
            Literal::Struct {
                export_name,
                fields,
                path,
            } => {
                if config.language == Language::C {
                    write!(out, "({})", export_name);
                } else {
                    write!(out, "{}", export_name);
                }

                write!(out, "{{ ");
                let mut is_first_field = true;
                // In C++, same order as defined is required.
                let ordered_fields = out.bindings().struct_field_names(path);
                for ordered_key in ordered_fields.iter() {
                    if let Some(ref lit) = fields.get(ordered_key) {
                        if !is_first_field {
                            write!(out, ", ");
                        } else {
                            is_first_field = false;
                        }
                        if config.language == Language::Cxx {
                            write!(out, "/* .{} = */ ", ordered_key);
                        } else {
                            write!(out, ".{} = ", ordered_key);
                        }
                        lit.write(config, out);
                    }
                }
                write!(out, " }}");
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct Constant {
    pub path: Path,
    pub export_name: String,
    pub ty: Type,
    pub value: Literal,
    pub cfg: Option<Cfg>,
    pub annotations: AnnotationSet,
    pub documentation: Documentation,
    pub associated_to: Option<Path>,
}

fn can_handle(ty: &Type, expr: &syn::Expr) -> bool {
    if ty.is_primitive_or_ptr_primitive() {
        return true;
    }
    match *expr {
        syn::Expr::Struct(_) => true,
        _ => false,
    }
}

impl Constant {
    pub fn load(
        path: Path,
        mod_cfg: Option<&Cfg>,
        ty: &syn::Type,
        expr: &syn::Expr,
        attrs: &[syn::Attribute],
        associated_to: Option<Path>,
    ) -> Result<Constant, String> {
        let ty = Type::load(ty)?;
        let mut ty = match ty {
            Some(ty) => ty,
            None => {
                return Err("Cannot have a zero sized const definition.".to_owned());
            }
        };

        if !can_handle(&ty, expr) {
            return Err("Unhandled const definition".to_owned());
        }

        let mut lit = Literal::load(&expr)?;

        if let Some(ref associated_to) = associated_to {
            ty.replace_self_with(associated_to);
            lit.replace_self_with(associated_to);
        }

        Ok(Constant::new(
            path,
            ty,
            lit,
            Cfg::append(mod_cfg, Cfg::load(attrs)),
            AnnotationSet::load(attrs)?,
            Documentation::load(attrs),
            associated_to,
        ))
    }

    pub fn new(
        path: Path,
        ty: Type,
        value: Literal,
        cfg: Option<Cfg>,
        annotations: AnnotationSet,
        documentation: Documentation,
        associated_to: Option<Path>,
    ) -> Self {
        let export_name = path.name().to_owned();
        Self {
            path,
            export_name,
            ty,
            value,
            cfg,
            annotations,
            documentation,
            associated_to,
        }
    }
}

impl Item for Constant {
    fn path(&self) -> &Path {
        &self.path
    }

    fn add_dependencies(&self, library: &Library, out: &mut Dependencies) {
        self.ty.add_dependencies(library, out);
    }

    fn export_name(&self) -> &str {
        &self.export_name
    }

    fn cfg(&self) -> Option<&Cfg> {
        self.cfg.as_ref()
    }

    fn annotations(&self) -> &AnnotationSet {
        &self.annotations
    }

    fn annotations_mut(&mut self) -> &mut AnnotationSet {
        &mut self.annotations
    }

    fn container(&self) -> ItemContainer {
        ItemContainer::Constant(self.clone())
    }

    fn rename_for_config(&mut self, config: &Config) {
        if self.associated_to.is_none() {
            config.export.rename(&mut self.export_name);
        }
        self.value.rename_for_config(config);
        self.ty.rename_for_config(config, &GenericParams::default()); // FIXME: should probably propagate something here
    }

    fn resolve_declaration_types(&mut self, resolver: &DeclarationTypeResolver) {
        self.ty.resolve_declaration_types(resolver);
    }
}

impl Constant {
    pub fn write_declaration<F: Write>(
        &self,
        config: &Config,
        out: &mut SourceWriter<F>,
        associated_to_struct: &Struct,
    ) {
        debug_assert!(self.associated_to.is_some());
        debug_assert!(config.language == Language::Cxx);
        debug_assert!(!associated_to_struct.is_transparent);
        debug_assert!(config.structure.associated_constants_in_body);
        debug_assert!(config.constant.allow_static_const);

        if let Type::ConstPtr(..) = self.ty {
            out.write("static ");
        } else {
            out.write("static const ");
        }
        self.ty.write(config, out);
        write!(out, " {};", self.export_name())
    }

    pub fn write<F: Write>(
        &self,
        config: &Config,
        out: &mut SourceWriter<F>,
        associated_to_struct: Option<&Struct>,
    ) {
        if let Some(assoc) = associated_to_struct {
            if assoc.is_generic() {
                return; // Not tested / implemented yet, so bail out.
            }
        }

        if !self.value.is_valid(out.bindings()) {
            return;
        }

        let associated_to_transparent = associated_to_struct.map_or(false, |s| s.is_transparent);

        let in_body = associated_to_struct.is_some()
            && config.language == Language::Cxx
            && config.structure.associated_constants_in_body
            && config.constant.allow_static_const
            && !associated_to_transparent;

        let condition = (&self.cfg).to_condition(config);
        condition.write_before(config, out);

        let name = if in_body {
            Cow::Owned(format!(
                "{}::{}",
                associated_to_struct.unwrap().export_name(),
                self.export_name(),
            ))
        } else if self.associated_to.is_none() {
            Cow::Borrowed(self.export_name())
        } else {
            let associated_name = match associated_to_struct {
                Some(s) => Cow::Borrowed(s.export_name()),
                None => {
                    let mut name = self.associated_to.as_ref().unwrap().name().to_owned();
                    config.export.rename(&mut name);
                    Cow::Owned(name)
                }
            };

            Cow::Owned(format!("{}_{}", associated_name, self.export_name()))
        };

        let value = match self.value {
            Literal::Struct {
                ref fields,
                ref path,
                ..
            } if out.bindings().struct_is_transparent(path) => &fields.iter().next().unwrap().1,
            _ => &self.value,
        };

        if config.constant.allow_static_const && config.language == Language::Cxx {
            out.write(if in_body { "inline " } else { "static " });
            if let Type::ConstPtr(..) = self.ty {
                // Nothing.
            } else {
                out.write("const ");
            }
            self.ty.write(config, out);
            write!(out, " {} = ", name);
            value.write(config, out);
            write!(out, ";");
        } else {
            write!(out, "#define {} ", name);
            value.write(config, out);
        }
        condition.write_after(config, out);
    }
}
<commit_msg>Handle all binary operations and make pattern match exhaustive<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/. */

use std::borrow::Cow;
use std::collections::HashMap;
use std::io::Write;

use syn::{self, UnOp};

use bindgen::config::{Config, Language};
use bindgen::declarationtyperesolver::DeclarationTypeResolver;
use bindgen::dependencies::Dependencies;
use bindgen::ir::{
    AnnotationSet, Cfg, ConditionWrite, Documentation, GenericParams, Item, ItemContainer, Path,
    Struct, ToCondition, Type,
};
use bindgen::library::Library;
use bindgen::writer::{Source, SourceWriter};
use bindgen::Bindings;

#[derive(Debug, Clone)]
pub enum Literal {
    Expr(String),
    PostfixUnaryOp {
        op: &'static str,
        value: Box<Literal>,
    },
    BinOp {
        left: Box<Literal>,
        op: &'static str,
        right: Box<Literal>,
    },
    Struct {
        path: Path,
        export_name: String,
        fields: HashMap<String, Literal>,
    },
}

impl Literal {
    fn replace_self_with(&mut self, self_ty: &Path) {
        match *self {
            Literal::PostfixUnaryOp { .. } | Literal::BinOp { .. } | Literal::Expr(..) => {}
            Literal::Struct {
                ref mut path,
                ref mut export_name,
                ref mut fields,
            } => {
                if path.replace_self_with(self_ty) {
                    *export_name = self_ty.name().to_owned();
                }
                for (ref _name, ref mut expr) in fields {
                    expr.replace_self_with(self_ty);
                }
            }
        }
    }

    fn is_valid(&self, bindings: &Bindings) -> bool {
        match *self {
            Literal::Expr(..) => true,
            Literal::PostfixUnaryOp { ref value, .. } => value.is_valid(bindings),
            Literal::BinOp {
                ref left,
                ref right,
                ..
            } => left.is_valid(bindings) && right.is_valid(bindings),
            Literal::Struct { ref path, .. } => bindings.struct_exists(path),
        }
    }
}

impl Literal {
    pub fn rename_for_config(&mut self, config: &Config) {
        match self {
            Literal::Struct {
                ref mut export_name,
                fields,
                ..
            } => {
                config.export.rename(export_name);
                for (_, lit) in fields {
                    lit.rename_for_config(config);
                }
            }
            Literal::PostfixUnaryOp { ref mut value, .. } => {
                value.rename_for_config(config);
            }
            Literal::BinOp {
                ref mut left,
                ref mut right,
                ..
            } => {
                left.rename_for_config(config);
                right.rename_for_config(config);
            }
            Literal::Expr(_) => {}
        }
    }

    // Translate from full blown `syn::Expr` into a simpler `Literal` type
    pub fn load(expr: &syn::Expr) -> Result<Literal, String> {
        match *expr {
            // Match binary expressions of the form `a * b`
            syn::Expr::Binary(ref bin_expr) => {
                let l = Self::load(&bin_expr.left)?;
                let r = Self::load(&bin_expr.right)?;
                let op = match bin_expr.op {
                    syn::BinOp::Add(..) => "+",
                    syn::BinOp::Sub(..) => "-",
                    syn::BinOp::Mul(..) => "*",
                    syn::BinOp::Div(..) => "/",
                    syn::BinOp::Rem(..) => "%",
                    syn::BinOp::And(..) => "&&",
                    syn::BinOp::Or(..) => "||",
                    syn::BinOp::BitXor(..) => "^",
                    syn::BinOp::BitAnd(..) => "&",
                    syn::BinOp::BitOr(..) => "|",
                    syn::BinOp::Shl(..) => "<<",
                    syn::BinOp::Shr(..) => ">>",
                    syn::BinOp::Eq(..) => "==",
                    syn::BinOp::Lt(..) => "<",
                    syn::BinOp::Le(..) => "<=",
                    syn::BinOp::Ne(..) => "!=",
                    syn::BinOp::Ge(..) => ">=",
                    syn::BinOp::Gt(..) => ">",
                    syn::BinOp::AddEq(..) => "+=",
                    syn::BinOp::SubEq(..) => "-=",
                    syn::BinOp::MulEq(..) => "*=",
                    syn::BinOp::DivEq(..) => "/=",
                    syn::BinOp::RemEq(..) => "%=",
                    syn::BinOp::BitXorEq(..) => "^=",
                    syn::BinOp::BitAndEq(..) => "&=",
                    syn::BinOp::BitOrEq(..) => "|=",
                    syn::BinOp::ShlEq(..) => ">>=",
                    syn::BinOp::ShrEq(..) => "<<=",
                };
                Ok(Literal::BinOp {
                    left: Box::new(l),
                    op,
                    right: Box::new(r),
                })
            }

            // Match literals like "one", 'a', 32 etc
            syn::Expr::Lit(syn::ExprLit { ref lit, .. }) => {
                match lit {
                    syn::Lit::Str(ref value) => {
                        Ok(Literal::Expr(format!("u8\"{}\"", value.value())))
                    }
                    syn::Lit::Byte(ref value) => Ok(Literal::Expr(format!("{}", value.value()))),
                    syn::Lit::Char(ref value) => Ok(Literal::Expr(match value.value() as u32 {
                        0..=255 => format!("'{}'", value.value().escape_default()),
                        other_code => format!(r"L'\u{:X}'", other_code),
                    })),
                    syn::Lit::Int(ref value) => {
                        Ok(Literal::Expr(value.base10_digits().to_string()))
                    }
                    syn::Lit::Float(ref value) => {
                        Ok(Literal::Expr(value.base10_digits().to_string()))
                    }
                    syn::Lit::Bool(ref value) => Ok(Literal::Expr(format!("{}", value.value))),
                    // TODO: Add support for byte string and Verbatim
                    _ => Err(format!("Unsupported literal expression. {:?}", *lit)),
                }
            }

            syn::Expr::Struct(syn::ExprStruct {
                ref path,
                ref fields,
                ..
            }) => {
                let struct_name = path.segments[0].ident.to_string();
                let mut field_map = HashMap::<String, Literal>::default();
                for field in fields {
                    let ident = match field.member {
                        syn::Member::Named(ref name) => name.to_string(),
                        syn::Member::Unnamed(ref index) => format!("_{}", index.index),
                    };
                    let key = ident.to_string();
                    let value = Literal::load(&field.expr)?;
                    field_map.insert(key, value);
                }
                Ok(Literal::Struct {
                    path: Path::new(struct_name.clone()),
                    export_name: struct_name,
                    fields: field_map,
                })
            }

            syn::Expr::Unary(syn::ExprUnary {
                ref op, ref expr, ..
            }) => match *op {
                UnOp::Neg(_) => {
                    let val = Self::load(expr)?;
                    Ok(Literal::PostfixUnaryOp {
                        op: "-",
                        value: Box::new(val),
                    })
                }
                _ => Err(format!("Unsupported Unary expression. {:?}", *op)),
            },

            _ => Err(format!("Unsupported expression. {:?}", *expr)),
        }
    }

    fn write<F: Write>(&self, config: &Config, out: &mut SourceWriter<F>) {
        match self {
            Literal::Expr(v) => write!(out, "{}", v),
            Literal::PostfixUnaryOp { op, ref value } => {
                write!(out, "{}", op);
                value.write(config, out);
            }
            Literal::BinOp {
                ref left,
                op,
                ref right,
            } => {
                left.write(config, out);
                write!(out, " {} ", op);
                right.write(config, out);
            }
            Literal::Struct {
                export_name,
                fields,
                path,
            } => {
                if config.language == Language::C {
                    write!(out, "({})", export_name);
                } else {
                    write!(out, "{}", export_name);
                }

                write!(out, "{{ ");
                let mut is_first_field = true;
                // In C++, same order as defined is required.
                let ordered_fields = out.bindings().struct_field_names(path);
                for ordered_key in ordered_fields.iter() {
                    if let Some(ref lit) = fields.get(ordered_key) {
                        if !is_first_field {
                            write!(out, ", ");
                        } else {
                            is_first_field = false;
                        }
                        if config.language == Language::Cxx {
                            write!(out, "/* .{} = */ ", ordered_key);
                        } else {
                            write!(out, ".{} = ", ordered_key);
                        }
                        lit.write(config, out);
                    }
                }
                write!(out, " }}");
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct Constant {
    pub path: Path,
    pub export_name: String,
    pub ty: Type,
    pub value: Literal,
    pub cfg: Option<Cfg>,
    pub annotations: AnnotationSet,
    pub documentation: Documentation,
    pub associated_to: Option<Path>,
}

fn can_handle(ty: &Type, expr: &syn::Expr) -> bool {
    if ty.is_primitive_or_ptr_primitive() {
        return true;
    }
    match *expr {
        syn::Expr::Struct(_) => true,
        _ => false,
    }
}

impl Constant {
    pub fn load(
        path: Path,
        mod_cfg: Option<&Cfg>,
        ty: &syn::Type,
        expr: &syn::Expr,
        attrs: &[syn::Attribute],
        associated_to: Option<Path>,
    ) -> Result<Constant, String> {
        let ty = Type::load(ty)?;
        let mut ty = match ty {
            Some(ty) => ty,
            None => {
                return Err("Cannot have a zero sized const definition.".to_owned());
            }
        };

        if !can_handle(&ty, expr) {
            return Err("Unhandled const definition".to_owned());
        }

        let mut lit = Literal::load(&expr)?;

        if let Some(ref associated_to) = associated_to {
            ty.replace_self_with(associated_to);
            lit.replace_self_with(associated_to);
        }

        Ok(Constant::new(
            path,
            ty,
            lit,
            Cfg::append(mod_cfg, Cfg::load(attrs)),
            AnnotationSet::load(attrs)?,
            Documentation::load(attrs),
            associated_to,
        ))
    }

    pub fn new(
        path: Path,
        ty: Type,
        value: Literal,
        cfg: Option<Cfg>,
        annotations: AnnotationSet,
        documentation: Documentation,
        associated_to: Option<Path>,
    ) -> Self {
        let export_name = path.name().to_owned();
        Self {
            path,
            export_name,
            ty,
            value,
            cfg,
            annotations,
            documentation,
            associated_to,
        }
    }
}

impl Item for Constant {
    fn path(&self) -> &Path {
        &self.path
    }

    fn add_dependencies(&self, library: &Library, out: &mut Dependencies) {
        self.ty.add_dependencies(library, out);
    }

    fn export_name(&self) -> &str {
        &self.export_name
    }

    fn cfg(&self) -> Option<&Cfg> {
        self.cfg.as_ref()
    }

    fn annotations(&self) -> &AnnotationSet {
        &self.annotations
    }

    fn annotations_mut(&mut self) -> &mut AnnotationSet {
        &mut self.annotations
    }

    fn container(&self) -> ItemContainer {
        ItemContainer::Constant(self.clone())
    }

    fn rename_for_config(&mut self, config: &Config) {
        if self.associated_to.is_none() {
            config.export.rename(&mut self.export_name);
        }
        self.value.rename_for_config(config);
        self.ty.rename_for_config(config, &GenericParams::default()); // FIXME: should probably propagate something here
    }

    fn resolve_declaration_types(&mut self, resolver: &DeclarationTypeResolver) {
        self.ty.resolve_declaration_types(resolver);
    }
}

impl Constant {
    pub fn write_declaration<F: Write>(
        &self,
        config: &Config,
        out: &mut SourceWriter<F>,
        associated_to_struct: &Struct,
    ) {
        debug_assert!(self.associated_to.is_some());
        debug_assert!(config.language == Language::Cxx);
        debug_assert!(!associated_to_struct.is_transparent);
        debug_assert!(config.structure.associated_constants_in_body);
        debug_assert!(config.constant.allow_static_const);

        if let Type::ConstPtr(..) = self.ty {
            out.write("static ");
        } else {
            out.write("static const ");
        }
        self.ty.write(config, out);
        write!(out, " {};", self.export_name())
    }

    pub fn write<F: Write>(
        &self,
        config: &Config,
        out: &mut SourceWriter<F>,
        associated_to_struct: Option<&Struct>,
    ) {
        if let Some(assoc) = associated_to_struct {
            if assoc.is_generic() {
                return; // Not tested / implemented yet, so bail out.
            }
        }

        if !self.value.is_valid(out.bindings()) {
            return;
        }

        let associated_to_transparent = associated_to_struct.map_or(false, |s| s.is_transparent);

        let in_body = associated_to_struct.is_some()
            && config.language == Language::Cxx
            && config.structure.associated_constants_in_body
            && config.constant.allow_static_const
            && !associated_to_transparent;

        let condition = (&self.cfg).to_condition(config);
        condition.write_before(config, out);

        let name = if in_body {
            Cow::Owned(format!(
                "{}::{}",
                associated_to_struct.unwrap().export_name(),
                self.export_name(),
            ))
        } else if self.associated_to.is_none() {
            Cow::Borrowed(self.export_name())
        } else {
            let associated_name = match associated_to_struct {
                Some(s) => Cow::Borrowed(s.export_name()),
                None => {
                    let mut name = self.associated_to.as_ref().unwrap().name().to_owned();
                    config.export.rename(&mut name);
                    Cow::Owned(name)
                }
            };

            Cow::Owned(format!("{}_{}", associated_name, self.export_name()))
        };

        let value = match self.value {
            Literal::Struct {
                ref fields,
                ref path,
                ..
            } if out.bindings().struct_is_transparent(path) => &fields.iter().next().unwrap().1,
            _ => &self.value,
        };

        if config.constant.allow_static_const && config.language == Language::Cxx {
            out.write(if in_body { "inline " } else { "static " });
            if let Type::ConstPtr(..) = self.ty {
                // Nothing.
            } else {
                out.write("const ");
            }
            self.ty.write(config, out);
            write!(out, " {} = ", name);
            value.write(config, out);
            write!(out, ";");
        } else {
            write!(out, "#define {} ", name);
            value.write(config, out);
        }
        condition.write_after(config, out);
    }
}
<|endoftext|>