<commit_before>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
mod conv;

use std::{
    error, fmt,
    io::{BufRead, BufReader},
    path::{Path, PathBuf},
    process::{Command, Stdio},
    time::Instant,
};

use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    CodeAction, CodeActionOrCommand, Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin,
    WorkDoneProgressEnd, WorkDoneProgressReport,
};

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckConfig {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    // XXX: drop order is significant
    cmd_send: Sender<CheckCommand>,
    handle: Option<jod_thread::JoinHandle<()>>,
    pub task_recv: Receiver<CheckTask>,
}

impl CheckWatcher {
    pub fn new(config: CheckConfig, workspace_root: PathBuf) -> CheckWatcher {
        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = jod_thread::spawn(move || {
            let mut check = CheckWatcherThread::new(config, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send, handle: Some(handle) }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        self.cmd_send.send(CheckCommand::Update).unwrap();
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { url: Url, diagnostic: Diagnostic, fixes: Vec<CodeActionOrCommand> },

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherThread {
    options: CheckConfig,
    workspace_root: PathBuf,
    watcher: WatchThread,
    last_update_req: Option<Instant>,
}

impl CheckWatcherThread {
    fn new(options: CheckConfig, workspace_root: PathBuf) -> CheckWatcherThread {
        CheckWatcherThread {
            options,
            workspace_root,
            watcher: WatchThread::dummy(),
            last_update_req: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        // If we rerun the thread, we need to discard the previous check results first
        self.clean_previous_results(task_send);

        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.watcher.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        std::mem::replace(&mut self.watcher.message_recv, never());
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req.take();
                task_send.send(CheckTask::ClearDiagnostics).unwrap();

                // Replace with a dummy watcher first so we drop the original and wait for completion
                std::mem::replace(&mut self.watcher, WatchThread::dummy());

                // Then create the actual new watcher
                self.watcher = WatchThread::new(&self.options, &self.workspace_root);
            }
        }
    }

    fn clean_previous_results(&self, task_send: &Sender<CheckTask>) {
        task_send.send(CheckTask::ClearDiagnostics).unwrap();
        task_send
            .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd { message: None })))
            .unwrap();
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result = map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root);
                if map_result.is_empty() {
                    return;
                }

                for MappedRustDiagnostic { location, diagnostic, fixes } in map_result {
                    let fixes = fixes
                        .into_iter()
                        .map(|fix| {
                            CodeAction { diagnostics: Some(vec![diagnostic.clone()]), ..fix }.into()
                        })
                        .collect();

                    task_send
                        .send(CheckTask::AddDiagnostic { url: location.uri, diagnostic, fixes })
                        .unwrap();
                }
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    fixes: Vec<CodeAction>,
}

/// WatchThread exists to wrap around the communication needed to be able to
/// run `cargo check` without blocking. Currently the Rust standard library
/// doesn't provide a way to read sub-process output without blocking, so we
/// have to wrap sub-processes output handling in a thread and pass messages
/// back over a channel.
/// The correct way to dispose of the thread is to drop it, on which the
/// sub-process will be killed, and the thread will be joined.
struct WatchThread {
    // XXX: drop order is significant
    message_recv: Receiver<CheckEvent>,
    _handle: Option<jod_thread::JoinHandle<()>>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

#[derive(Debug)]
pub struct CargoError(String);

impl fmt::Display for CargoError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Cargo failed: {}", self.0)
    }
}
impl error::Error for CargoError {}

fn run_cargo(
    args: &[String],
    current_dir: Option<&Path>,
    on_message: &mut dyn FnMut(cargo_metadata::Message) -> bool,
) -> Result<(), CargoError> {
    let mut command = Command::new("cargo");
    if let Some(current_dir) = current_dir {
        command.current_dir(current_dir);
    }

    let mut child = command
        .args(args)
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .stdin(Stdio::null())
        .spawn()
        .expect("couldn't launch cargo");

    // We manually read a line at a time, instead of using serde's
    // stream deserializers, because the deserializer cannot recover
    // from an error, resulting in it getting stuck, because we try to
    // be resillient against failures.
    //
    // Because cargo only outputs one JSON object per line, we can
    // simply skip a line if it doesn't parse, which just ignores any
    // erroneus output.
    let stdout = BufReader::new(child.stdout.take().unwrap());
    let mut read_at_least_one_message = false;

    for line in stdout.lines() {
        let line = match line {
            Ok(line) => line,
            Err(err) => {
                log::error!("Couldn't read line from cargo: {}", err);
                continue;
            }
        };

        let message = serde_json::from_str::<cargo_metadata::Message>(&line);
        let message = match message {
            Ok(message) => message,
            Err(err) => {
                log::error!("Invalid json from cargo check, ignoring ({}): {:?} ", err, line);
                continue;
            }
        };

        read_at_least_one_message = true;

        if !on_message(message) {
            break;
        }
    }

    // It is okay to ignore the result, as it only errors if the process is already dead
    let _ = child.kill();

    let err_msg = match child.wait() {
        Ok(exit_code) if !exit_code.success() && !read_at_least_one_message => {
            // FIXME: Read the stderr to display the reason, see `read2()` reference in PR comment:
            // https://github.com/rust-analyzer/rust-analyzer/pull/3632#discussion_r395605298
            format!(
                "the command produced no valid metadata (exit code: {:?}): cargo {}",
                exit_code,
                args.join(" ")
            )
        }
        Err(err) => format!("io error: {:?}", err),
        Ok(_) => return Ok(()),
    };

    Err(CargoError(err_msg))
}

impl WatchThread {
    fn dummy() -> WatchThread {
        WatchThread { message_recv: never(), _handle: None }
    }

    fn new(options: &CheckConfig, workspace_root: &Path) -> WatchThread {
        let mut args: Vec<String> = vec![
            options.command.clone(),
            "--workspace".to_string(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", workspace_root.display()),
        ];
        if options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let workspace_root = workspace_root.to_owned();
        let handle = if options.enable {
            Some(jod_thread::spawn(move || {
                // If we trigger an error here, we will do so in the loop instead,
                // which will break out of the loop, and continue the shutdown
                let _ = message_send.send(CheckEvent::Begin);

                let res = run_cargo(&args, Some(&workspace_root), &mut |message| {
                    // Skip certain kinds of messages to only spend time on what's useful
                    match &message {
                        Message::CompilerArtifact(artifact) if artifact.fresh => return true,
                        Message::BuildScriptExecuted(_) => return true,
                        Message::Unknown => return true,
                        _ => {}
                    }

                    // if the send channel was closed, we want to shutdown
                    message_send.send(CheckEvent::Msg(message)).is_ok()
                });

                if let Err(err) = res {
                    // FIXME: make the `message_send` to be `Sender<Result<CheckEvent, CargoError>>`
                    // to display user-caused misconfiguration errors instead of just logging them here
                    log::error!("Cargo watcher failed {:?}", err);
                }

                // We can ignore any error here, as we are already in the progress
                // of shutting down.
                let _ = message_send.send(CheckEvent::End);
            }))
        } else {
            None
        };
        WatchThread { message_recv, _handle: handle }
    }
}
<commit_msg>Streamline flycheck implementation<commit_after>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
mod conv;

use std::{
    error, fmt,
    io::{BufRead, BufReader},
    path::{Path, PathBuf},
    process::{Command, Stdio},
    time::Instant,
};

use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    CodeAction, CodeActionOrCommand, Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin,
    WorkDoneProgressEnd, WorkDoneProgressReport,
};

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckConfig {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    // XXX: drop order is significant
    cmd_send: Sender<CheckCommand>,
    handle: Option<jod_thread::JoinHandle<()>>,
    pub task_recv: Receiver<CheckTask>,
}

impl CheckWatcher {
    pub fn new(config: CheckConfig, workspace_root: PathBuf) -> CheckWatcher {
        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = jod_thread::spawn(move || {
            let mut check = CheckWatcherThread::new(config, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send, handle: Some(handle) }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        self.cmd_send.send(CheckCommand::Update).unwrap();
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { url: Url, diagnostic: Diagnostic, fixes: Vec<CodeActionOrCommand> },

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherThread {
    options: CheckConfig,
    workspace_root: PathBuf,
    last_update_req: Option<Instant>,
    // XXX: drop order is significant
    message_recv: Receiver<CheckEvent>,
    /// WatchThread exists to wrap around the communication needed to be able to
    /// run `cargo check` without blocking. Currently the Rust standard library
    /// doesn't provide a way to read sub-process output without blocking, so we
    /// have to wrap sub-processes output handling in a thread and pass messages
    /// back over a channel.
    check_process: Option<jod_thread::JoinHandle<()>>,
}

impl CheckWatcherThread {
    fn new(options: CheckConfig, workspace_root: PathBuf) -> CheckWatcherThread {
        CheckWatcherThread {
            options,
            workspace_root,
            last_update_req: None,
            message_recv: never(),
            check_process: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        // If we rerun the thread, we need to discard the previous check results first
        self.clean_previous_results(task_send);

        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        self.message_recv = never();
                        self.check_process = None;
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req = None;
                task_send.send(CheckTask::ClearDiagnostics).unwrap();
                self.restart_check_process();
            }
        }
    }

    fn clean_previous_results(&self, task_send: &Sender<CheckTask>) {
        task_send.send(CheckTask::ClearDiagnostics).unwrap();
        task_send
            .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd { message: None })))
            .unwrap();
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result = map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root);
                if map_result.is_empty() {
                    return;
                }

                for MappedRustDiagnostic { location, diagnostic, fixes } in map_result {
                    let fixes = fixes
                        .into_iter()
                        .map(|fix| {
                            CodeAction { diagnostics: Some(vec![diagnostic.clone()]), ..fix }.into()
                        })
                        .collect();

                    task_send
                        .send(CheckTask::AddDiagnostic { url: location.uri, diagnostic, fixes })
                        .unwrap();
                }
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }

    fn restart_check_process(&mut self) {
        // First, clear and cancel the old thread
        self.message_recv = never();
        self.check_process = None;
        if !self.options.enable {
            return;
        }

        let mut args: Vec<String> = vec![
            self.options.command.clone(),
            "--workspace".to_string(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", self.workspace_root.display()),
        ];
        if self.options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(self.options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let workspace_root = self.workspace_root.to_owned();
        self.message_recv = message_recv;
        self.check_process = Some(jod_thread::spawn(move || {
            // If we trigger an error here, we will do so in the loop instead,
            // which will break out of the loop, and continue the shutdown
            let _ = message_send.send(CheckEvent::Begin);

            let res = run_cargo(&args, Some(&workspace_root), &mut |message| {
                // Skip certain kinds of messages to only spend time on what's useful
                match &message {
                    Message::CompilerArtifact(artifact) if artifact.fresh => return true,
                    Message::BuildScriptExecuted(_) => return true,
                    Message::Unknown => return true,
                    _ => {}
                }

                // if the send channel was closed, we want to shutdown
                message_send.send(CheckEvent::Msg(message)).is_ok()
            });

            if let Err(err) = res {
                // FIXME: make the `message_send` to be `Sender<Result<CheckEvent, CargoError>>`
                // to display user-caused misconfiguration errors instead of just logging them here
                log::error!("Cargo watcher failed {:?}", err);
            }

            // We can ignore any error here, as we are already in the progress
            // of shutting down.
            let _ = message_send.send(CheckEvent::End);
        }))
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    fixes: Vec<CodeAction>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

#[derive(Debug)]
pub struct CargoError(String);

impl fmt::Display for CargoError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Cargo failed: {}", self.0)
    }
}
impl error::Error for CargoError {}

fn run_cargo(
    args: &[String],
    current_dir: Option<&Path>,
    on_message: &mut dyn FnMut(cargo_metadata::Message) -> bool,
) -> Result<(), CargoError> {
    let mut command = Command::new("cargo");
    if let Some(current_dir) = current_dir {
        command.current_dir(current_dir);
    }

    let mut child = command
        .args(args)
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .stdin(Stdio::null())
        .spawn()
        .expect("couldn't launch cargo");

    // We manually read a line at a time, instead of using serde's
    // stream deserializers, because the deserializer cannot recover
    // from an error, resulting in it getting stuck, because we try to
    // be resillient against failures.
    //
    // Because cargo only outputs one JSON object per line, we can
    // simply skip a line if it doesn't parse, which just ignores any
    // erroneus output.
    let stdout = BufReader::new(child.stdout.take().unwrap());
    let mut read_at_least_one_message = false;

    for line in stdout.lines() {
        let line = match line {
            Ok(line) => line,
            Err(err) => {
                log::error!("Couldn't read line from cargo: {}", err);
                continue;
            }
        };

        let message = serde_json::from_str::<cargo_metadata::Message>(&line);
        let message = match message {
            Ok(message) => message,
            Err(err) => {
                log::error!("Invalid json from cargo check, ignoring ({}): {:?} ", err, line);
                continue;
            }
        };

        read_at_least_one_message = true;

        if !on_message(message) {
            break;
        }
    }

    // It is okay to ignore the result, as it only errors if the process is already dead
    let _ = child.kill();

    let err_msg = match child.wait() {
        Ok(exit_code) if !exit_code.success() && !read_at_least_one_message => {
            // FIXME: Read the stderr to display the reason, see `read2()` reference in PR comment:
            // https://github.com/rust-analyzer/rust-analyzer/pull/3632#discussion_r395605298
            format!(
                "the command produced no valid metadata (exit code: {:?}): cargo {}",
                exit_code,
                args.join(" ")
            )
        }
        Err(err) => format!("io error: {:?}", err),
        Ok(_) => return Ok(()),
    };

    Err(CargoError(err_msg))
}
<|endoftext|>


<commit_before>trait Parser<T, E> {
    fn next(&mut self) -> Option<T>;
    fn preview(&self) -> Option<&T>;
    fn current_pos(&self) -> (i32, i32);
    fn error<S: Into<String>>(&self, message: S) -> E;
}

macro_rules! next {
    ($p:expr) => {
        match $p.next() {
            Some(x) => x,
            None => return Err($p.error("unexpected eof")),
        }
    }
}

macro_rules! predicate {
    ($p:expr, $pred:expr) => {{
        let x = next!($p);
        if $pred(x) {
            x
        } else {
            return Err($p.error(format!("unexpected token {}", x)));
        }
    }}
}

macro_rules! atom {
    ($p:expr, $expected:expr) => {{
        let x = next!($p);
        if x == $expected {
            x
        } else {
            return Err(
                $p.error(format!("unexpected token {}, expected {}", x, $expected))
            );
        }
    }}
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::vec_deque::VecDeque;

    struct TP {
        input: VecDeque<i32>,
    }

    impl TP {
        fn new(input: &[i32]) -> TP {
            TP { input: VecDeque::from(input.to_vec()) }
        }
    }

    impl Parser<i32, String> for TP {
        fn next(&mut self) -> Option<i32> {
            self.input.pop_front()
        }

        fn preview(&self) -> Option<&i32> {
            self.input.front()
        }

        fn current_pos(&self) -> (i32, i32) {
            (0, 0)
        }

        fn error<S: Into<String>>(&self, message: S) -> String {
            message.into()
        }
    }

    macro_rules! with_res {
        (test $i:ident, $b:block) => {
            #[test]
            fn $i() {
                let result = || {
                    $b;
                    Ok(())
                };
                result().unwrap()
            }
        };
        (panic $e:expr, $i:ident, $b:block) => {
            #[test]
            #[should_panic(expected = $e)]
            fn $i() {
                let result = || {
                    $b;
                    Ok(())
                };
                result().unwrap()
            }
        };
    }

    with_res!(test next_success, {
        let mut p = TP::new(&[1, 2, 3]);
        assert_eq!(next!(p), 1);
        assert_eq!(next!(p), 2);
        assert_eq!(next!(p), 3);
    });

    with_res!(panic "unexpected eof", next_fail_empty, {
        let mut p = TP::new(&[]);
        assert_eq!(next!(p), 1);
    });

    with_res!(test predicate_success, {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 2);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 4);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 6);
    });

    with_res!(panic "unexpected eof", predicate_fail_empty, {
        let mut p = TP::new(&[]);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 2);
    });

    with_res!(panic "unexpected token 3", predicate_fail_not_satisfy, {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 3);
    });

    with_res!(test atom_success, {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(atom!(p, 2), 2);
        assert_eq!(atom!(p, 4), 4);
        assert_eq!(atom!(p, 6), 6);
    });

    with_res!(panic "unexpected eof", atom_fail_empty, {
        let mut p = TP::new(&[]);
        assert_eq!(atom!(p, 2), 2);
    });

    with_res!(panic "unexpected token 5, expected 4", atom_fail_not_expected, {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(atom!(p, 3), 3);
        assert_eq!(atom!(p, 4), 4);
    });
}
<commit_msg>Implement string combinator<commit_after>trait Parser<T, E> {
    fn next(&mut self) -> Option<T>;
    fn preview(&self) -> Option<&T>;
    fn current_pos(&self) -> (i32, i32);
    fn error<S: Into<String>>(&self, message: S) -> E;
}

macro_rules! next {
    ($p:expr) => {
        match $p.next() {
            Some(x) => x,
            None => return Err($p.error("unexpected eof")),
        }
    }
}

macro_rules! predicate {
    ($p:expr, $pred:expr) => {{
        let x = next!($p);
        if $pred(x) {
            x
        } else {
            return Err($p.error(format!("unexpected token {}", x)));
        }
    }}
}

macro_rules! atom {
    ($p:expr, $expected:expr) => {{
        let x = next!($p);
        if x == $expected {
            x
        } else {
            return Err(
                $p.error(format!("unexpected token {}, expected {}", x, $expected))
            );
        }
    }}
}

macro_rules! string {
    ($p:expr, $iter:expr, $t:ident) => {{
        let i = $iter;
        for c in &i {
            atom!($p, *c);
        }
        $t::from_iter(i)
    }}
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::vec_deque::VecDeque;
    use std::iter::FromIterator;

    struct TP {
        input: VecDeque<i32>,
    }

    impl TP {
        fn new(input: &[i32]) -> TP {
            TP { input: VecDeque::from(input.to_vec()) }
        }
    }

    impl Parser<i32, String> for TP {
        fn next(&mut self) -> Option<i32> {
            self.input.pop_front()
        }

        fn preview(&self) -> Option<&i32> {
            self.input.front()
        }

        fn current_pos(&self) -> (i32, i32) {
            (0, 0)
        }

        fn error<S: Into<String>>(&self, message: S) -> String {
            message.into()
        }
    }

    macro_rules! with_res {
        (test $i:ident, $b:block) => {
            #[test]
            fn $i() {
                let result = || {
                    $b;
                    Ok(())
                };
                result().unwrap()
            }
        };
        (panic $e:expr, $i:ident, $b:block) => {
            #[test]
            #[should_panic(expected = $e)]
            fn $i() {
                let result = || {
                    $b;
                    Ok(())
                };
                result().unwrap()
            }
        };
    }

    with_res!(test next_success, {
        let mut p = TP::new(&[1, 2, 3]);
        assert_eq!(next!(p), 1);
        assert_eq!(next!(p), 2);
        assert_eq!(next!(p), 3);
    });

    with_res!(panic "unexpected eof", next_fail_empty, {
        let mut p = TP::new(&[]);
        assert_eq!(next!(p), 1);
    });

    with_res!(test predicate_success, {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 2);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 4);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 6);
    });

    with_res!(panic "unexpected eof", predicate_fail_empty, {
        let mut p = TP::new(&[]);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 2);
    });

    with_res!(panic "unexpected token 3", predicate_fail_not_satisfy, {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(predicate!(p, |x| x % 2 == 0), 3);
    });

    with_res!(test atom_success, {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(atom!(p, 2), 2);
        assert_eq!(atom!(p, 4), 4);
        assert_eq!(atom!(p, 6), 6);
    });

    with_res!(panic "unexpected eof", atom_fail_empty, {
        let mut p = TP::new(&[]);
        assert_eq!(atom!(p, 2), 2);
    });

    with_res!(panic "unexpected token 5, expected 4", atom_fail_not_expected, {
        let mut p = TP::new(&[3, 5, 7]);
        assert_eq!(atom!(p, 3), 3);
        assert_eq!(atom!(p, 4), 4);
    });

    with_res!(test string_success, {
        let mut p = TP::new(&[2, 4, 6]);
        assert_eq!(string!(p, vec![2, 4, 6], Vec), vec![2, 4, 6]);
    });

    with_res!(panic "unexpected eof", string_fail_empty, {
        let mut p = TP::new(&[]);
        assert_eq!(string!(p, vec![2, 4, 6], Vec), [2, 4, 6]);
    });

    with_res!(panic "unexpected token 5, expected 4", string_fail_not_expected, {
        let mut p = TP::new(&[2, 5, 6]);
        assert_eq!(string!(p, vec![2, 4, 6], Vec), vec![2, 4, 6]);
    });
}
<|endoftext|>


<commit_before>use std::{fs, io, path::Path};

#[allow(unused_imports)]
use common::{
    eyre::{bail, eyre, Result},
    tracing,
};
use path_utils::lexiclean::Lexiclean;

#[cfg(feature = "tar")]
pub use ::tar;

#[cfg(feature = "zip")]
pub use ::zip;

#[cfg(feature = "flate2")]
pub use ::flate2;

#[cfg(feature = "xz2")]
pub use ::xz2;

#[cfg(feature = "zstd")]
pub use ::zstd;

/// Extract an archive to a destination
#[allow(unused_variables)]
pub fn extract(archive: &Path, dest: &Path, strip: usize, subdir: Option<&str>) -> Result<()> {
    tracing::info!("Extracting `{}` to `{}`", archive.display(), dest.display());

    let ext = match archive.extension() {
        Some(ext) => ext.to_string_lossy().to_string(),
        None => {
            bail!(
                "Archive file name `{}` has no extension so I don't know how to extract from it",
                archive.display()
            );
        }
    };

    match ext.as_str() {
        #[cfg(feature = "zip")]
        "zip" => extract_zip(archive, dest, strip, subdir),
        #[cfg(feature = "tar")]
        _ => extract_tar(&ext, archive, dest, strip, subdir),
        #[cfg(not(feature = "tar"))]
        _ => bail!("Unhandled archive extension {}", ext),
    }
}

/// Extract a tar archive
#[cfg(feature = "tar")]
pub fn extract_tar(
    ext: &str,
    archive: &Path,
    dest: &Path,
    strip: usize,
    subdir: Option<&str>,
) -> Result<()> {
    let file = fs::File::open(&archive)?;
    let mut archive = tar::Archive::new(match ext {
        "tar" => Box::new(file) as Box<dyn io::Read>,
        #[cfg(feature = "tar-gz")]
        "gz" | "tgz" => Box::new(flate2::read::GzDecoder::new(file)),
        #[cfg(feature = "tar-xz")]
        "xz" => Box::new(xz2::read::XzDecoder::new(file)),
        #[cfg(feature = "tar-zstd")]
        "zst" | "zstd" => Box::new(zstd::stream::Decoder::new(file)),
        _ => bail!("Unhandled archive extension {}", ext),
    });

    let subdir: Vec<String> = match subdir {
        Some(subdir) => subdir.split('/').map(String::from).collect(),
        None => Vec::new(),
    };

    let mut count = 0;
    for mut entry in archive.entries()?.flatten() {
        let path = entry.path()?.display().to_string();
        let mut components: Vec<String> = path.split('/').map(String::from).collect();

        // Strip path components
        if strip > 0 {
            components.drain(0..strip);
        }

        // Only extract entries in specified subdirectories
        let mut extract = true;
        for dir in &subdir {
            if components.first() == Some(dir) {
                components.remove(0);
            } else {
                extract = false;
                break;
            }
        }
        if !extract {
            continue;
        }

        // Do not extract paths that traverse out of destination
        let out_path = dest.join(&components.join("/")).lexiclean();
        if out_path.strip_prefix(dest).is_err() {
            continue;
        }

        // Unpack to destination
        entry.unpack(&out_path)?;
        count += 1;
    }

    tracing::debug!("Extracted {} entries", count);
    Ok(())
}

/// Extract a zip archive
#[cfg(feature = "zip")]
pub fn extract_zip(archive: &Path, dest: &Path, strip: usize, subdir: Option<&str>) -> Result<()> {
    let file = fs::File::open(&archive)?;
    let mut archive = zip::ZipArchive::new(file)?;

    let subdir: Vec<String> = match subdir {
        Some(subdir) => subdir.split('/').map(String::from).collect(),
        None => Vec::new(),
    };

    let mut count = 0;
    for index in 0..archive.len() {
        let mut entry = archive.by_index(index)?;
        let path = entry
            .enclosed_name()
            .ok_or_else(|| eyre!("Zip archive entry has no enclosed name"))?
            .display()
            .to_string();
        let mut components: Vec<String> = path.split('/').map(String::from).collect();

        // Strip path components
        if strip > 0 {
            components.drain(0..strip);
        }

        // Only extract entries in specified subdirectories
        let mut extract = true;
        for dir in &subdir {
            if components.first() == Some(dir) {
                components.remove(0);
            } else {
                extract = false;
                break;
            }
        }
        if !extract {
            continue;
        }

        // Do not extract paths that traverse out of destination
        let out_path = dest.join(&components.join("/")).lexiclean();
        if out_path.strip_prefix(dest).is_err() {
            continue;
        }

        // Unpack to destination
        if entry.is_file() {
            if let Some(parent) = out_path.parent() {
                fs::create_dir_all(parent)?;
            }
            let mut out_file = fs::File::create(out_path)?;
            std::io::copy(&mut entry, &mut out_file)?;
            count += 1;
        }
    }

    tracing::debug!("Extracted {} entries", count);
    Ok(())
}
<commit_msg>fix(Archive utilities): Ensure parent directories are present<commit_after>use std::{fs, io, path::Path};

#[allow(unused_imports)]
use common::{
    eyre::{bail, eyre, Result},
    tracing,
};
use path_utils::lexiclean::Lexiclean;

#[cfg(feature = "tar")]
pub use ::tar;

#[cfg(feature = "zip")]
pub use ::zip;

#[cfg(feature = "flate2")]
pub use ::flate2;

#[cfg(feature = "xz2")]
pub use ::xz2;

#[cfg(feature = "zstd")]
pub use ::zstd;

/// Extract an archive to a destination
#[allow(unused_variables)]
pub fn extract(archive: &Path, dest: &Path, strip: usize, subdir: Option<&str>) -> Result<()> {
    tracing::info!("Extracting `{}` to `{}`", archive.display(), dest.display());

    let ext = match archive.extension() {
        Some(ext) => ext.to_string_lossy().to_string(),
        None => {
            bail!(
                "Archive file name `{}` has no extension so I don't know how to extract from it",
                archive.display()
            );
        }
    };

    match ext.as_str() {
        #[cfg(feature = "zip")]
        "zip" => extract_zip(archive, dest, strip, subdir),
        #[cfg(feature = "tar")]
        _ => extract_tar(&ext, archive, dest, strip, subdir),
        #[cfg(not(feature = "tar"))]
        _ => bail!("Unhandled archive extension {}", ext),
    }
}

/// Extract a tar archive
#[cfg(feature = "tar")]
pub fn extract_tar(
    ext: &str,
    archive: &Path,
    dest: &Path,
    strip: usize,
    subdir: Option<&str>,
) -> Result<()> {
    use std::fs::create_dir_all;

    let file = fs::File::open(&archive)?;
    let mut archive = tar::Archive::new(match ext {
        "tar" => Box::new(file) as Box<dyn io::Read>,
        #[cfg(feature = "tar-gz")]
        "gz" | "tgz" => Box::new(flate2::read::GzDecoder::new(file)),
        #[cfg(feature = "tar-xz")]
        "xz" => Box::new(xz2::read::XzDecoder::new(file)),
        #[cfg(feature = "tar-zstd")]
        "zst" | "zstd" => Box::new(zstd::stream::Decoder::new(file)),
        _ => bail!("Unhandled archive extension {}", ext),
    });

    let subdir: Vec<String> = match subdir {
        Some(subdir) => subdir.split('/').map(String::from).collect(),
        None => Vec::new(),
    };

    let mut count = 0;
    for mut entry in archive.entries()?.flatten() {
        let path = entry.path()?.display().to_string();
        let mut components: Vec<String> = path.split('/').map(String::from).collect();

        // Strip path components
        if strip > 0 {
            components.drain(0..strip);
        }

        // Only extract entries in specified subdirectories
        let mut extract = true;
        for dir in &subdir {
            if components.first() == Some(dir) {
                components.remove(0);
            } else {
                extract = false;
                break;
            }
        }
        if !extract {
            continue;
        }

        // Do not extract paths that traverse out of destination
        let out_path = dest.join(&components.join("/")).lexiclean();
        if out_path.strip_prefix(dest).is_err() {
            continue;
        }

        // Ensure the parent directories of the output path exist
        if let Some(parent) = out_path.parent() {
            create_dir_all(&parent)?;
        }

        // Unpack to destination
        entry.unpack(&out_path)?;
        count += 1;
    }

    tracing::debug!("Extracted {} entries", count);
    Ok(())
}

/// Extract a zip archive
#[cfg(feature = "zip")]
pub fn extract_zip(archive: &Path, dest: &Path, strip: usize, subdir: Option<&str>) -> Result<()> {
    let file = fs::File::open(&archive)?;
    let mut archive = zip::ZipArchive::new(file)?;

    let subdir: Vec<String> = match subdir {
        Some(subdir) => subdir.split('/').map(String::from).collect(),
        None => Vec::new(),
    };

    let mut count = 0;
    for index in 0..archive.len() {
        let mut entry = archive.by_index(index)?;
        let path = entry
            .enclosed_name()
            .ok_or_else(|| eyre!("Zip archive entry has no enclosed name"))?
            .display()
            .to_string();
        let mut components: Vec<String> = path.split('/').map(String::from).collect();

        // Strip path components
        if strip > 0 {
            components.drain(0..strip);
        }

        // Only extract entries in specified subdirectories
        let mut extract = true;
        for dir in &subdir {
            if components.first() == Some(dir) {
                components.remove(0);
            } else {
                extract = false;
                break;
            }
        }
        if !extract {
            continue;
        }

        // Do not extract paths that traverse out of destination
        let out_path = dest.join(&components.join("/")).lexiclean();
        if out_path.strip_prefix(dest).is_err() {
            continue;
        }

        // Unpack to destination
        if entry.is_file() {
            if let Some(parent) = out_path.parent() {
                fs::create_dir_all(parent)?;
            }
            let mut out_file = fs::File::create(out_path)?;
            std::io::copy(&mut entry, &mut out_file)?;
            count += 1;
        }
    }

    tracing::debug!("Extracted {} entries", count);
    Ok(())
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

extern crate arrow;
extern crate criterion;
extern crate rand;

use std::mem::size_of;

use criterion::*;
use rand::distributions::Standard;
use rand::{thread_rng, Rng};

use arrow::builder::*;

// Build arrays with 512k elements.
const BATCH_SIZE: usize = 8 << 10;
const NUM_BATCHES: usize = 64;

fn bench_primitive(c: &mut Criterion) {
    let data: [i64; BATCH_SIZE] = [100; BATCH_SIZE];
    c.bench(
        "bench_primitive",
        Benchmark::new("bench_primitive", move |b| {
            b.iter(|| {
                let mut builder = Int64Builder::new(64);
                for _ in 0..NUM_BATCHES {
                    let _ = black_box(builder.push_slice(&data[..]));
                }
                black_box(builder.finish());
            })
        })
        .throughput(Throughput::Bytes(
            (data.len() * NUM_BATCHES * size_of::<i64>()) as u32,
        )),
    );
}

fn bench_bool(c: &mut Criterion) {
    let data: Vec<bool> = thread_rng()
        .sample_iter(&Standard)
        .take(BATCH_SIZE)
        .collect();
    let data_len = data.len();
    c.bench(
        "bench_bool",
        Benchmark::new("bench_bool", move |b| {
            b.iter(|| {
                let mut builder = BooleanBuilder::new(64);
                for _ in 0..NUM_BATCHES {
                    let _ = black_box(builder.push_slice(&data[..]));
                }
                black_box(builder.finish());
            })
        })
        .throughput(Throughput::Bytes(
            (data_len * NUM_BATCHES * size_of::<bool>()) as u32,
        )),
    );
}

criterion_group!(benches, bench_primitive, bench_bool);
criterion_main!(benches);
<commit_msg>ARROW-4282: [Rust] builder benchmark is broken<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

extern crate arrow;
extern crate criterion;
extern crate rand;

use std::mem::size_of;

use criterion::*;
use rand::distributions::Standard;
use rand::{thread_rng, Rng};

use arrow::builder::*;

// Build arrays with 512k elements.
const BATCH_SIZE: usize = 8 << 10;
const NUM_BATCHES: usize = 64;

fn bench_primitive(c: &mut Criterion) {
    let data: [i64; BATCH_SIZE] = [100; BATCH_SIZE];
    c.bench(
        "bench_primitive",
        Benchmark::new("bench_primitive", move |b| {
            b.iter(|| {
                let mut builder = Int64Builder::new(64);
                for _ in 0..NUM_BATCHES {
                    let _ = black_box(builder.append_slice(&data[..]));
                }
                black_box(builder.finish());
            })
        })
        .throughput(Throughput::Bytes(
            (data.len() * NUM_BATCHES * size_of::<i64>()) as u32,
        )),
    );
}

fn bench_bool(c: &mut Criterion) {
    let data: Vec<bool> = thread_rng()
        .sample_iter(&Standard)
        .take(BATCH_SIZE)
        .collect();
    let data_len = data.len();
    c.bench(
        "bench_bool",
        Benchmark::new("bench_bool", move |b| {
            b.iter(|| {
                let mut builder = BooleanBuilder::new(64);
                for _ in 0..NUM_BATCHES {
                    let _ = black_box(builder.append_slice(&data[..]));
                }
                black_box(builder.finish());
            })
        })
        .throughput(Throughput::Bytes(
            (data_len * NUM_BATCHES * size_of::<bool>()) as u32,
        )),
    );
}

criterion_group!(benches, bench_primitive, bench_bool);
criterion_main!(benches);
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::io::{self, Write};
use std::net::Shutdown;
use std::rc::Rc;
use mio::Token;
use mio::net::TcpStream;
use mio::{Event, PollOpt, Ready};

use super::binary;
use super::close_listener::CloseListener;
use super::ipv4_packet::MAX_PACKET_LENGTH;
use super::ipv4_packet_buffer::IPv4PacketBuffer;
use super::router::Router;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;

const TAG: &'static str = "Client";

pub struct Client {
    id: u32,
    stream: TcpStream,
    client_to_network: IPv4PacketBuffer,
    network_to_client: StreamBuffer,
    router: Router,
    closed: bool,
    close_listener: Box<CloseListener<Client>>,
    token: Token,
    // number of remaining bytes of "id" to send to the client before relaying any data
    pending_id_bytes: usize,
}

impl Client {
    pub fn new(id: u32, selector: &mut Selector, stream: TcpStream, close_listener: Box<CloseListener<Client>>) -> io::Result<Rc<RefCell<Self>>> {
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            stream: stream,
            client_to_network: IPv4PacketBuffer::new(),
            network_to_client: StreamBuffer::new(16 * MAX_PACKET_LENGTH),
            router: Router::new(),
            closed: false,
            close_listener: close_listener,
            token: Token(0), // default value, will be set afterwards
            pending_id_bytes: 4,
        }));
        // set client as router owner
        rc.borrow_mut().router.set_client(Rc::downgrade(&rc));

        let rc_clone = rc.clone();
        let handler = move |selector: &mut Selector, ready| {
            let mut self_ref = rc_clone.borrow_mut();
            self_ref.on_ready(selector, ready);
        };
        {
            let mut self_ref = rc.borrow_mut();
            // on start, we are interested only in writing (we must first send the client id)
            let token = selector.register(&self_ref.stream, handler, Ready::writable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    pub fn id(&self) -> u32 {
        return self.id;
    }

    pub fn router(&mut self) -> &mut Router {
        &mut self.router
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        selector.deregister(&self.stream, self.token);
        // shutdown only (there is no close), the socket will be closed on drop
        self.stream.shutdown(Shutdown::Both);
        self.router.clear();
        self.close_listener.on_closed(self);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if self.must_send_id() {
            match self.send_id() {
                Ok(_) => {
                    if self.pending_id_bytes == 0 {
                        debug!("Client id #{} sent to client", self.id);
                    }
                }
                Err(_) => {
                    error!(target: TAG, "Cannot write client id #{}", self.id);
                    self.close(selector);
                }
            }
        } else {
            match self.write() {
                Ok(_) => self.process_pending(),
                Err(_) => {
                    error!(target: TAG, "Cannot write");
                    self.close(selector);
                }
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        match self.read() {
            Ok(_) => self.push_to_network(selector),
            Err(_) => {
                error!(target: TAG, "Cannot read");
                self.close(selector);
            }
        }
    }

    fn send_id(&mut self) -> io::Result<()> {
        assert!(self.must_send_id());
        let raw_id = binary::to_byte_array(self.id);
        let w = self.stream.write(&raw_id[4 - self.pending_id_bytes..])?;
        self.pending_id_bytes -= w;
        Ok(())
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let ready = if self.network_to_client.is_empty() {
            Ready::readable()
        } else {
            Ready::readable() | Ready::writable()
        };
        selector.reregister(&self.stream, self.token, ready, PollOpt::level())
    }

    fn read(&mut self) -> io::Result<()> {
        self.client_to_network.read_from(&mut self.stream)
    }

    fn write(&mut self) -> io::Result<()> {
        self.network_to_client.write_to(&mut self.stream)?;
        Ok(())
    }

    fn push_to_network(&mut self, selector: &mut Selector) {
        while self.push_one_packet_to_network(selector) {
            self.client_to_network.next();
        }
    }

    fn push_one_packet_to_network(&mut self, selector: &mut Selector) -> bool {
        match self.client_to_network.as_ipv4_packet() {
            Some(ref packet) => {
                self.router.send_to_network(packet);
                true
            }
            None => false
        }
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn process_pending(&mut self) {
        // TODO
    }

    fn clean_expired_connections() {
        // TODO
    }

    fn must_send_id(&self) -> bool{
        self.pending_id_bytes > 0
    }
}
<commit_msg>Improve scoping<commit_after>use std::cell::RefCell;
use std::io::{self, Write};
use std::net::Shutdown;
use std::rc::Rc;
use mio::Token;
use mio::net::TcpStream;
use mio::{Event, PollOpt, Ready};

use super::binary;
use super::close_listener::CloseListener;
use super::ipv4_packet::MAX_PACKET_LENGTH;
use super::ipv4_packet_buffer::IPv4PacketBuffer;
use super::router::Router;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;

const TAG: &'static str = "Client";

pub struct Client {
    id: u32,
    stream: TcpStream,
    client_to_network: IPv4PacketBuffer,
    network_to_client: StreamBuffer,
    router: Router,
    closed: bool,
    close_listener: Box<CloseListener<Client>>,
    token: Token,
    // number of remaining bytes of "id" to send to the client before relaying any data
    pending_id_bytes: usize,
}

impl Client {
    pub fn new(id: u32, selector: &mut Selector, stream: TcpStream, close_listener: Box<CloseListener<Client>>) -> io::Result<Rc<RefCell<Self>>> {
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            stream: stream,
            client_to_network: IPv4PacketBuffer::new(),
            network_to_client: StreamBuffer::new(16 * MAX_PACKET_LENGTH),
            router: Router::new(),
            closed: false,
            close_listener: close_listener,
            token: Token(0), // default value, will be set afterwards
            pending_id_bytes: 4,
        }));
        // set client as router owner
        rc.borrow_mut().router.set_client(Rc::downgrade(&rc));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            // on start, we are interested only in writing (we must first send the client id)
            let token = selector.register(&self_ref.stream, handler, Ready::writable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    pub fn id(&self) -> u32 {
        return self.id;
    }

    pub fn router(&mut self) -> &mut Router {
        &mut self.router
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        selector.deregister(&self.stream, self.token);
        // shutdown only (there is no close), the socket will be closed on drop
        self.stream.shutdown(Shutdown::Both);
        self.router.clear();
        self.close_listener.on_closed(self);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if self.must_send_id() {
            match self.send_id() {
                Ok(_) => {
                    if self.pending_id_bytes == 0 {
                        debug!("Client id #{} sent to client", self.id);
                    }
                }
                Err(_) => {
                    error!(target: TAG, "Cannot write client id #{}", self.id);
                    self.close(selector);
                }
            }
        } else {
            match self.write() {
                Ok(_) => self.process_pending(),
                Err(_) => {
                    error!(target: TAG, "Cannot write");
                    self.close(selector);
                }
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        match self.read() {
            Ok(_) => self.push_to_network(selector),
            Err(_) => {
                error!(target: TAG, "Cannot read");
                self.close(selector);
            }
        }
    }

    fn send_id(&mut self) -> io::Result<()> {
        assert!(self.must_send_id());
        let raw_id = binary::to_byte_array(self.id);
        let w = self.stream.write(&raw_id[4 - self.pending_id_bytes..])?;
        self.pending_id_bytes -= w;
        Ok(())
    }

    fn update_interests(&mut self, selector: &mut Selector) -> io::Result<()> {
        let ready = if self.network_to_client.is_empty() {
            Ready::readable()
        } else {
            Ready::readable() | Ready::writable()
        };
        selector.reregister(&self.stream, self.token, ready, PollOpt::level())
    }

    fn read(&mut self) -> io::Result<()> {
        self.client_to_network.read_from(&mut self.stream)
    }

    fn write(&mut self) -> io::Result<()> {
        self.network_to_client.write_to(&mut self.stream)?;
        Ok(())
    }

    fn push_to_network(&mut self, selector: &mut Selector) {
        while self.push_one_packet_to_network(selector) {
            self.client_to_network.next();
        }
    }

    fn push_one_packet_to_network(&mut self, selector: &mut Selector) -> bool {
        match self.client_to_network.as_ipv4_packet() {
            Some(ref packet) => {
                self.router.send_to_network(packet);
                true
            }
            None => false
        }
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn process_pending(&mut self) {
        // TODO
    }

    fn clean_expired_connections() {
        // TODO
    }

    fn must_send_id(&self) -> bool{
        self.pending_id_bytes > 0
    }
}
<|endoftext|>


<commit_before>//! An rope.
//!
//! A rope is an efficient data structure for large mutable strings. It's
//! essentially a binary tree whose leaves are strings.
//!
//! For more information, see the following resources:
//! + http://scienceblogs.com/goodmath/2009/01/26/ropes-twining-together-strings/
//! + https://www.ibm.com/developerworks/library/j-ropes/
//! + http://citeseer.ist.psu.edu/viewdoc/download?doi=10.1.1.14.9450&rep=rep1&type=pdf

#![feature(const_fn)]

use std::ops;
use std::convert;

#[derive(Debug)]
pub struct Rope {
    // can we get away with having these be of &str or will they need
    // to be string?
    root: Node
}

impl Rope {

    /// Returns a new empty Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::new();
    /// assert_eq!(an_rope.length(), 0);
    /// ```
    pub const fn new() -> Rope {
        Rope { root: Node::None }
    }

    /// Returns the length of this Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_empty_rope = Rope::new();
    /// assert_eq!(an_empty_rope.length(), 0);
    /// ```
    /// ```
    /// use an_rope::Rope;
    /// let mut an_empty_rope = Rope::from(String::from(""));
    /// assert_eq!(an_empty_rope.length(), 0);
    /// ```
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::from(String::from("a string"));
    /// assert_eq!(an_rope.length(), "a string".len());
    /// ```
    pub fn length(&self) -> usize {
        unimplemented!()
    }

    /// Appends a rope to the end of this Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::from(String::from("abcd"));
    /// an_rope.append(Rope::from(String::from("efgh")))
    /// assert_eq!(an_rope, Rope::from(String::from("abcdefgh")));
    /// ```
    pub fn append(&mut self, other: Rope) {
        unimplemented!()
    }

    /// Prepends a rope to the front of this Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::from(String::from("efgh"));
    /// an_rope.preend(Rope::from(String::from("abcd")))
    /// assert_eq!(an_rope, Rope::from(String::from("abcdefgh")));
    /// ```
    pub fn prepend(&mut self, other: Rope) {
        unimplemented!()
    }
}

#[derive(Debug)]
enum Node { /// A leaf node
            // todo: is box<str> the right choice?
            Leaf(Box<str>)
          , /// A branch node
            Branch { l: Box<Node>, r: Box<Node> }
          , /// Nothing
            None
}

impl Node {

    /// Computes the weight of a node
    fn weight(&self) -> usize {
        unimplemented!()
    }
}

impl convert::From<String> for Rope {
    fn from(string: String) -> Rope {
        unimplemented!()
    }
}

impl ops::Add for Rope {
    type Output = Rope;

    // This is the concat operation.
    fn add(self, other: Rope) -> Rope {
        unimplemented!()
    }
}

impl ops::Index<usize> for Rope {
    type Output = char;

    // Index a character
    fn index(&self, i: usize) -> &char {
        unimplemented!()
    }
}

// slicing operators -----------------------------------------
impl ops::Index<ops::Range<usize>> for Rope {
    type Output = str;

    // Index a substring
    fn index(&self, i: ops::Range<usize>) -> &str {
        unimplemented!()
    }
}

impl ops::Index<ops::RangeTo<usize>> for Rope {
    type Output = str;

    fn index(&self, i: ops::RangeTo<usize>) -> &str {
        unimplemented!()
    }
}

impl ops::Index<ops::RangeFrom<usize>> for Rope {
    type Output = str;

    fn index(&self, i: ops::RangeFrom<usize>) -> &str {
        unimplemented!()
    }
}

impl ops::IndexMut<ops::Range<usize>> for Rope {
    fn index_mut(&mut self, i: ops::Range<usize>) -> &mut str {
        unimplemented!()
    }
}

impl ops::IndexMut<ops::RangeTo<usize>> for Rope {
    fn index_mut(&mut self, i: ops::RangeTo<usize>) -> &mut str {
        unimplemented!()
    }
}

impl ops::IndexMut<ops::RangeFrom<usize>> for Rope {
    fn index_mut(&mut self, i: ops::RangeFrom<usize>) -> &mut str {
        unimplemented!()
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}
<commit_msg>rope: more unimplemented API<commit_after>//! An rope.
//!
//! A rope is an efficient data structure for large mutable strings. It's
//! essentially a binary tree whose leaves are strings.
//!
//! For more information, see the following resources:
//! + http://scienceblogs.com/goodmath/2009/01/26/ropes-twining-together-strings/
//! + https://www.ibm.com/developerworks/library/j-ropes/
//! + http://citeseer.ist.psu.edu/viewdoc/download?doi=10.1.1.14.9450&rep=rep1&type=pdf

#![feature(const_fn)]

use std::cmp;
use std::ops;
use std::convert;

#[derive(Debug)]
pub struct Rope {
    // can we get away with having these be of &str or will they need
    // to be string?
    root: Node
}

impl Rope {

    /// Returns a new empty Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::new();
    /// assert_eq!(an_rope.length(), 0);
    /// ```
    pub const fn new() -> Rope {
        Rope { root: Node::None }
    }

    /// Returns the length of this Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_empty_rope = Rope::new();
    /// assert_eq!(an_empty_rope.length(), 0);
    /// ```
    /// ```
    /// use an_rope::Rope;
    /// let mut an_empty_rope = Rope::from(String::from(""));
    /// assert_eq!(an_empty_rope.length(), 0);
    /// ```
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::from(String::from("a string"));
    /// assert_eq!(an_rope.length(), "a string".len());
    /// ```
    pub fn length(&self) -> usize {
        unimplemented!()
    }

    /// Appends a rope to the end of this Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::from(String::from("abcd"));
    /// an_rope.append(Rope::from(String::from("efgh")))
    /// assert_eq!(an_rope, Rope::from(String::from("abcdefgh")));
    /// ```
    pub fn append(&mut self, other: Rope) {
        unimplemented!()
    }

    /// Prepends a rope to the front of this Rope
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::from(String::from("efgh"));
    /// an_rope.preend(Rope::from(String::from("abcd")))
    /// assert_eq!(an_rope, Rope::from(String::from("abcdefgh")));
    /// ```
    pub fn prepend(&mut self, other: Rope) {
        unimplemented!()
    }

    /// Splits the rope into two ropes at the given index.
    ///
    /// Consumes this rope.
    ///
    /// # Examples
    /// ```
    /// use an_rope::Rope;
    /// let mut an_rope = Rope::from(String::from("abcd"));
    /// let (ab, cd) = an_rope.split(2);
    /// assert_eq!(ab, Rope::from(String::from("ab")));
    /// assert_eq!(cd, Rope::from(String::from("cd")));
    /// ```
    pub fn split(self, index: usize) -> (Rope, Rope) {
        unimplemented!()
    }
}

#[derive(Debug)]
enum Node { /// A leaf node
            // todo: is box<str> the right choice?
            Leaf(Box<str>)
          , /// A branch node
            Branch { l: Box<Node>, r: Box<Node> }
          , /// Nothing
            None
}

impl Node {

    /// Computes the weight of a node
    fn weight(&self) -> usize {
        unimplemented!()
    }
}

impl convert::From<String> for Rope {
    fn from(string: String) -> Rope {
        unimplemented!()
    }
}

//-- comparisons ----------------------------------------------------
impl cmp::PartialEq for Rope {
    fn eq(&self, other: &Rope) -> bool {
        unimplemented!()
    }
}

impl cmp::PartialEq<str> for Rope {
    fn eq(&self, other: &str) -> bool {
        unimplemented!()
    }
}

//-- concatenation --------------------------------------------------
impl ops::Add for Rope {
    type Output = Rope;

    // This is the concat operation.
    fn add(self, other: Rope) -> Rope {
        unimplemented!()
    }
}

impl ops::Index<usize> for Rope {
    type Output = char;

    // Index a character
    fn index(&self, i: usize) -> &char {
        unimplemented!()
    }
}

//-- slicing operators ----------------------------------------------
impl ops::Index<ops::Range<usize>> for Rope {
    type Output = str;

    // Index a substring
    fn index(&self, i: ops::Range<usize>) -> &str {
        unimplemented!()
    }
}

impl ops::Index<ops::RangeTo<usize>> for Rope {
    type Output = str;

    fn index(&self, i: ops::RangeTo<usize>) -> &str {
        unimplemented!()
    }
}

impl ops::Index<ops::RangeFrom<usize>> for Rope {
    type Output = str;

    fn index(&self, i: ops::RangeFrom<usize>) -> &str {
        unimplemented!()
    }
}

impl ops::IndexMut<ops::Range<usize>> for Rope {
    fn index_mut(&mut self, i: ops::Range<usize>) -> &mut str {
        unimplemented!()
    }
}

impl ops::IndexMut<ops::RangeTo<usize>> for Rope {
    fn index_mut(&mut self, i: ops::RangeTo<usize>) -> &mut str {
        unimplemented!()
    }
}

impl ops::IndexMut<ops::RangeFrom<usize>> for Rope {
    fn index_mut(&mut self, i: ops::RangeFrom<usize>) -> &mut str {
        unimplemented!()
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}
<|endoftext|>


<commit_before>use kata_io as io;

use consts::*;
use frame::*;
use proto::*;

const SUBPACKET_SIZE: usize = 1024 * 8;
const SUBPACKET_PER_ACK: usize = 10;

#[derive(Debug, PartialEq)]
enum State {
    /// Waiting ZRINIT invite (do nothing)
    WaitingInit,

    /// Sending ZRQINIT
    SendingZRQINIT,

    /// Sending ZFILE frame
    SendingZFILE,

    /// Do nothing, just waiting for ZPOS
    WaitingZPOS,

    /// Sending ZDATA & subpackets
    SendingData,

    /// Sending ZFIN
    SendingZFIN,

    /// All works done, exiting
    Done,
}

impl State {
    fn new() -> State {
        State::WaitingInit
    }

    fn next(self, frame: &Frame) -> State {
        match (self, frame.get_frame_type()) {
            (State::WaitingInit, ZRINIT) => State::SendingZFILE,
            (State::WaitingInit, _) => State::SendingZRQINIT,

            (State::SendingZRQINIT, ZRINIT) => State::SendingZFILE,

            (State::SendingZFILE, ZRPOS) => State::SendingData,
            (State::SendingZFILE, ZRINIT) => State::WaitingZPOS,

            (State::WaitingZPOS, ZRPOS) => State::SendingData,

            (State::SendingData, ZACK) => State::SendingData,
            (State::SendingData, ZRPOS) => State::SendingData,
            (State::SendingData, ZRINIT) => State::SendingZFIN,

            (State::SendingZFIN, ZFIN) => State::Done,

            (s, _) => {
                error!("Unexpected (state, frame) combination: {:#?} {}", s, frame);
                s // don't change current state
            }
        }
    }
}

pub fn send<CI, CO, DI>(
    mut channel_in: CI,
    mut channel_out: CO,
    mut r: DI,
    filename: &str,
    filesize: Option<u32>,
) -> io::Result<()>
where
    CI: io::Read,
    CO: io::Write,
    DI: io::Read + io::Seek,
{
    let mut data = [0; SUBPACKET_SIZE];
    let mut offset: u32;

    write_zrqinit(&mut channel_out)?;

    let mut state = State::new();

    while state != State::Done {
        channel_out.flush()?;

        if !find_zpad(&mut channel_in)? {
            continue;
        }

        let frame = match parse_header(&mut channel_in)? {
            Some(x) => x,
            None => {
                write_znak(&mut channel_out)?;
                continue;
            }
        };

        state = state.next(&frame);
        debug!("State: {:?}", state);

        // do things according new state
        match state {
            State::SendingZRQINIT => {
                write_zrqinit(&mut channel_out)?;
            }
            State::SendingZFILE => {
                write_zfile(&mut channel_out, filename, filesize)?;
            }
            State::SendingData => {
                offset = frame.get_count();
                r.seek(io::SeekFrom::Start(offset as u64))?;

                let num = r.read(&mut data)?;

                if num == 0 {
                    write_zeof(&mut channel_out, offset)?;
                } else {
                    // ZBIN32|ZDATA
                    // ZCRCG - best perf
                    // ZCRCQ - mid perf
                    // ZCRCW - worst perf
                    // ZCRCE - send at end
                    write_zdata(&mut channel_out, offset)?;

                    let mut i = 0;
                    loop {
                        i += 1;

                        write_zlde_data(&mut channel_out, ZCRCG, &data[..num])?;
                        offset += num as u32;

                        let num = r.read(&mut data)?;
                        if num < data.len() || i >= SUBPACKET_PER_ACK {
                            write_zlde_data(&mut channel_out, ZCRCW, &data[..num])?;
                            break;
                        }
                    }
                }
            }
            State::SendingZFIN => {
                write_zfin(&mut channel_out)?;
            }
            State::Done => {
                write_over_and_out(&mut channel_out)?;
            }
            _ => (),
        }
    }

    Ok(())
}
<commit_msg>ZMODEM: sz argument readability rename<commit_after>use kata_io as io;

use consts::*;
use frame::*;
use proto::*;

const SUBPACKET_SIZE: usize = 1024 * 8;
const SUBPACKET_PER_ACK: usize = 10;

#[derive(Debug, PartialEq)]
enum State {
    /// Waiting ZRINIT invite (do nothing)
    WaitingInit,

    /// Sending ZRQINIT
    SendingZRQINIT,

    /// Sending ZFILE frame
    SendingZFILE,

    /// Do nothing, just waiting for ZPOS
    WaitingZPOS,

    /// Sending ZDATA & subpackets
    SendingData,

    /// Sending ZFIN
    SendingZFIN,

    /// All works done, exiting
    Done,
}

impl State {
    fn new() -> State {
        State::WaitingInit
    }

    fn next(self, frame: &Frame) -> State {
        match (self, frame.get_frame_type()) {
            (State::WaitingInit, ZRINIT) => State::SendingZFILE,
            (State::WaitingInit, _) => State::SendingZRQINIT,

            (State::SendingZRQINIT, ZRINIT) => State::SendingZFILE,

            (State::SendingZFILE, ZRPOS) => State::SendingData,
            (State::SendingZFILE, ZRINIT) => State::WaitingZPOS,

            (State::WaitingZPOS, ZRPOS) => State::SendingData,

            (State::SendingData, ZACK) => State::SendingData,
            (State::SendingData, ZRPOS) => State::SendingData,
            (State::SendingData, ZRINIT) => State::SendingZFIN,

            (State::SendingZFIN, ZFIN) => State::Done,

            (s, _) => {
                error!("Unexpected (state, frame) combination: {:#?} {}", s, frame);
                s // don't change current state
            }
        }
    }
}

pub fn send<CI, CO, DI>(
    mut channel_in: CI,
    mut channel_out: CO,
    mut data_in: DI,
    filename: &str,
    filesize: Option<u32>,
) -> io::Result<()>
where
    CI: io::Read,
    CO: io::Write,
    DI: io::Read + io::Seek,
{
    let mut data = [0; SUBPACKET_SIZE];
    let mut offset: u32;

    write_zrqinit(&mut channel_out)?;

    let mut state = State::new();

    while state != State::Done {
        channel_out.flush()?;

        if !find_zpad(&mut channel_in)? {
            continue;
        }

        let frame = match parse_header(&mut channel_in)? {
            Some(x) => x,
            None => {
                write_znak(&mut channel_out)?;
                continue;
            }
        };

        state = state.next(&frame);
        debug!("State: {:?}", state);

        // do things according new state
        match state {
            State::SendingZRQINIT => {
                write_zrqinit(&mut channel_out)?;
            }
            State::SendingZFILE => {
                write_zfile(&mut channel_out, filename, filesize)?;
            }
            State::SendingData => {
                offset = frame.get_count();
                data_in.seek(io::SeekFrom::Start(offset as u64))?;

                let num = data_in.read(&mut data)?;

                if num == 0 {
                    write_zeof(&mut channel_out, offset)?;
                } else {
                    // ZBIN32|ZDATA
                    // ZCRCG - best perf
                    // ZCRCQ - mid perf
                    // ZCRCW - worst perf
                    // ZCRCE - send at end
                    write_zdata(&mut channel_out, offset)?;

                    let mut i = 0;
                    loop {
                        i += 1;

                        write_zlde_data(&mut channel_out, ZCRCG, &data[..num])?;
                        offset += num as u32;

                        let num = data_in.read(&mut data)?;
                        if num < data.len() || i >= SUBPACKET_PER_ACK {
                            write_zlde_data(&mut channel_out, ZCRCW, &data[..num])?;
                            break;
                        }
                    }
                }
            }
            State::SendingZFIN => {
                write_zfin(&mut channel_out)?;
            }
            State::Done => {
                write_over_and_out(&mut channel_out)?;
            }
            _ => (),
        }
    }

    Ok(())
}
<|endoftext|>