<commit_before>// Copyright 2018 Google Inc.
//
// Licensed 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
//
//     https://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.

use colored::*;
use connection::Connection;
use control_file;
use evcxr;
use evcxr::CommandContext;
use failure::Error;
use json;
use json::JsonValue;
use jupyter_message::JupyterMessage;
use std;
use std::collections::HashMap;
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use std::time;
use zmq;

// Note, to avoid potential deadlocks, each thread should lock at most one mutex at a time.
#[derive(Clone)]
pub(crate) struct Server {
    iopub: Arc<Mutex<Connection>>,
    _stdin: Arc<Mutex<Connection>>,
    latest_execution_request: Arc<Mutex<Option<JupyterMessage>>>,
    shutdown_requested_receiver: Arc<Mutex<mpsc::Receiver<()>>>,
    shutdown_requested_sender: Arc<Mutex<mpsc::Sender<()>>>,
}

impl Server {
    pub(crate) fn start(config: &control_file::Control) -> Result<Server, Error> {
        use zmq::SocketType;

        let zmq_context = zmq::Context::new();
        let heartbeat = bind_socket(config, config.hb_port, zmq_context.socket(SocketType::REP)?)?;
        let shell_socket = bind_socket(
            config,
            config.shell_port,
            zmq_context.socket(SocketType::ROUTER)?,
        )?;
        let control_socket = bind_socket(
            config,
            config.control_port,
            zmq_context.socket(SocketType::ROUTER)?,
        )?;
        let stdin_socket = bind_socket(
            config,
            config.stdin_port,
            zmq_context.socket(SocketType::ROUTER)?,
        )?;
        let iopub = Arc::new(Mutex::new(bind_socket(
            config,
            config.iopub_port,
            zmq_context.socket(SocketType::PUB)?,
        )?));

        let (shutdown_requested_sender, shutdown_requested_receiver) = mpsc::channel();

        let server = Server {
            iopub,
            latest_execution_request: Arc::new(Mutex::new(None)),
            _stdin: Arc::new(Mutex::new(stdin_socket)),
            shutdown_requested_receiver: Arc::new(Mutex::new(shutdown_requested_receiver)),
            shutdown_requested_sender: Arc::new(Mutex::new(shutdown_requested_sender)),
        };

        let (execution_sender, execution_receiver) = mpsc::channel();
        let (execution_response_sender, execution_response_receiver) = mpsc::channel();

        thread::spawn(move || Self::handle_hb(&heartbeat));
        server.start_thread(move |server: Server| server.handle_control(control_socket));
        server.start_thread(move |server: Server| {
            server.handle_shell(
                shell_socket,
                &execution_sender,
                &execution_response_receiver,
            )
        });
        let (context, outputs) = CommandContext::new()?;
        server.start_thread(move |server: Server| {
            server.handle_execution_requests(
                context,
                &execution_receiver,
                &execution_response_sender,
            )
        });
        server
            .clone()
            .start_output_pass_through_thread("stdout", outputs.stdout);
        server
            .clone()
            .start_output_pass_through_thread("stderr", outputs.stderr);
        Ok(server)
    }

    pub(crate) fn wait_for_shutdown(&self) {
        self.shutdown_requested_receiver
            .lock()
            .unwrap()
            .recv()
            .unwrap();
    }

    fn signal_shutdown(&self) {
        self.shutdown_requested_sender
            .lock()
            .unwrap()
            .send(())
            .unwrap();
    }

    fn start_thread<F>(&self, body: F)
    where
        F: FnOnce(Server) -> Result<(), Error> + std::marker::Send + 'static,
    {
        let server_clone = self.clone();
        thread::spawn(|| {
            if let Err(error) = body(server_clone) {
                eprintln!("{:?}", error);
            }
        });
    }

    fn handle_hb(connection: &Connection) -> Result<(), Error> {
        let mut message = zmq::Message::new()?;
        loop {
            connection.socket.recv(&mut message, 0)?;
            connection.socket.send(b"ping", 0)?;
        }
    }

    fn handle_execution_requests(
        self,
        mut context: CommandContext,
        receiver: &mpsc::Receiver<JupyterMessage>,
        execution_reply_sender: &mpsc::Sender<JupyterMessage>,
    ) -> Result<(), Error> {
        let mut execution_count = 1;
        loop {
            let message = receiver.recv()?;

            // If we want this clone to be cheaper, we probably only need the header, not the
            // whole message.
            *self.latest_execution_request.lock().unwrap() = Some(message.clone());
            let code = message.code();
            execution_count += 1;
            message
                .new_message("execute_input")
                .with_content(object! {
                    "execution_count" => execution_count,
                    "code" => code
                })
                .send(&mut *self.iopub.lock().unwrap())?;
            let reply = match context.execute(code) {
                Ok(output) => {
                    if !output.is_empty() {
                        // Increase the odds that stdout will have been finished being sent. A
                        // less hacky alternative would be to add a print statement, then block
                        // waiting for it.
                        thread::sleep(time::Duration::from_millis(1));
                        let mut data = HashMap::new();
                        // At the time of writing the json crate appears to have a generic From
                        // implementation for a Vec<T> where T implements Into<JsonValue>. It also
                        // has conversion from HashMap<String, JsonValue>, but it doesn't have
                        // conversion from HashMap<String, T>. Perhaps send a PR? For now, we
                        // convert the values manually.
                        for (k, v) in output.content_by_mime_type {
                            if k.contains("json") {
                                data.insert(k, json::parse(&v).unwrap_or_else(|_| json::from(v)));
                            } else {
                                data.insert(k, json::from(v));
                            }
                        }
                        message
                            .new_message("execute_result")
                            .with_content(object! {
                                "execution_count" => execution_count,
                                "data" => data,
                                "metadata" => HashMap::new(),
                            })
                            .send(&mut *self.iopub.lock().unwrap())?;
                    }
                    if let Some(duration) = output.timing {
                        // TODO replace by duration.as_millis() when stable
                        let ms =
                            duration.as_secs() * 1000 + u64::from(duration.subsec_millis());
                        let mut data = HashMap::new();
                        data.insert(
                            "text/html".into(),
                            json::from(format!(
                                "Took {}ms",
                                ms
                            )),
                        );
                        message
                            .new_message("execute_result")
                            .with_content(object! {
                                "execution_count" => execution_count,
                                "data" => data,
                                "metadata" => HashMap::new(),
                            })
                            .send(&mut *self.iopub.lock().unwrap())?;
                    }
                    message.new_reply().with_content(object! {
                        "status" => "ok",
                        "execution_count" => execution_count
                    })
                }
                Err(errors) => {
                    self.emit_errors(&errors, &message)?;
                    message.new_reply().with_content(object! {
                        "status" => "error",
                        "execution_count" => execution_count,
                    })
                }
            };
            execution_reply_sender.send(reply)?;
        }
    }

    fn handle_shell(
        self,
        mut connection: Connection,
        execution_channel: &mpsc::Sender<JupyterMessage>,
        execution_reply_receiver: &mpsc::Receiver<JupyterMessage>,
    ) -> Result<(), Error> {
        loop {
            let message = JupyterMessage::read(&mut connection)?;
            // Processing of every message should be enclosed between "busy" and "idle"
            // see https://jupyter-client.readthedocs.io/en/latest/messaging.html#messages-on-the-shell-router-dealer-channel
            // Jupiter Lab doesn't use the kernel until it received "idle" for kernel_info_request
            message
                .new_message("status")
                .with_content(object! {"execution_state" => "busy"})
                .send(&mut *self.iopub.lock().unwrap())?;
            let idle = message
                .new_message("status")
                .with_content(object! {"execution_state" => "idle"});
            if message.message_type() == "kernel_info_request" {
                message
                    .new_reply()
                    .with_content(kernel_info())
                    .send(&mut connection)?;
            } else if message.message_type() == "is_complete_request" {
                message
                    .new_reply()
                    .with_content(object! {"status" => "complete"})
                    .send(&mut connection)?;
            } else if message.message_type() == "execute_request" {
                execution_channel.send(message)?;
                execution_reply_receiver.recv()?.send(&mut connection)?;
            } else if message.message_type() == "comm_open" {
                message
                    .new_message("comm_close")
                    .with_content(message.get_content().clone())
                    .send(&mut connection)?;
            } else {
                eprintln!(
                    "Got unrecognized message type on shell channel: {}",
                    message.message_type()
                );
            }
            idle.send(&mut *self.iopub.lock().unwrap())?;
        }
    }

    fn handle_control(self, mut connection: Connection) -> Result<(), Error> {
        loop {
            let message = JupyterMessage::read(&mut connection)?;
            match message.message_type() {
                "shutdown_request" => self.signal_shutdown(),
                "interrupt_request" => {
                    message.new_reply().send(&mut connection)?;
                    eprintln!(
                        "Rust doesn't support interrupting execution. Perhaps restart kernel?"
                    );
                }
                _ => {
                    eprintln!(
                        "Got unrecognized message type on control channel: {}",
                        message.message_type()
                    );
                }
            }
        }
    }

    fn start_output_pass_through_thread(
        self,
        output_name: &'static str,
        channel: mpsc::Receiver<String>,
    ) {
        thread::spawn(move || {
            while let Ok(line) = channel.recv() {
                let mut message = None;
                if let Some(exec_request) = &*self.latest_execution_request.lock().unwrap() {
                    message = Some(exec_request.new_message("stream"));
                }
                if let Some(message) = message {
                    if let Err(error) = message
                        .with_content(object! {
                            "name" => output_name,
                            "text" => format!("{}\n", line),
                        })
                        .send(&mut *self.iopub.lock().unwrap())
                    {
                        eprintln!("{}", error);
                    }
                }
            }
        });
    }

    fn emit_errors(
        &self,
        errors: &evcxr::Error,
        parent_message: &JupyterMessage,
    ) -> Result<(), Error> {
        match errors {
            evcxr::Error::CompilationErrors(errors) => {
                for error in errors {
                    let message = format!("{}", error.message().bright_red());
                    if error.is_from_user_code() {
                        let mut traceback = Vec::new();
                        for spanned_message in error.spanned_messages() {
                            for line in &spanned_message.lines {
                                traceback.push(line.clone());
                            }
                            if let Some(span) = &spanned_message.span {
                                let mut carrots = String::new();
                                for _ in 1..span.start_column {
                                    carrots.push(' ');
                                }
                                for _ in span.start_column..span.end_column {
                                    carrots.push('^');
                                }
                                traceback.push(format!(
                                    "{} {}",
                                    carrots.bright_red(),
                                    spanned_message.label.bright_blue()
                                ));
                            } else {
                                traceback.push(spanned_message.label.clone());
                            }
                        }
                        traceback.push(error.message());
                        parent_message
                            .new_message("error")
                            .with_content(object! {
                                "ename" => "Error",
                                "evalue" => error.message(),
                                "traceback" => traceback,
                            })
                            .send(&mut *self.iopub.lock().unwrap())?;
                    } else {
                        parent_message
                            .new_message("error")
                            .with_content(object! {
                                "ename" => "Error",
                                "evalue" => error.message(),
                                "traceback" => array![
                                    message
                                ],
                            })
                            .send(&mut *self.iopub.lock().unwrap())?;
                    }
                }
            }
            error => {
                let displayed_error = format!("{}", error);
                parent_message
                    .new_message("error")
                    .with_content(object! {
                        "ename" => "Error",
                        "evalue" => displayed_error.clone(),
                        "traceback" => array![displayed_error],
                    })
                    .send(&mut *self.iopub.lock().unwrap())?;
            }
        }
        Ok(())
    }
}

fn bind_socket(
    config: &control_file::Control,
    port: u16,
    socket: zmq::Socket,
) -> Result<Connection, Error> {
    let endpoint = format!("{}://{}:{}", config.transport, config.ip, port);
    socket.bind(&endpoint)?;
    Ok(Connection::new(socket, &config.key)?)
}

/// See [Kernel info documentation](https://jupyter-client.readthedocs.io/en/stable/messaging.html#kernel-info)
fn kernel_info() -> JsonValue {
    object! {
        "protocol_version" => "5.3",
        "implementation" => env!("CARGO_PKG_NAME"),
        "implementation_version" => env!("CARGO_PKG_VERSION"),
        "language_info" => object!{
            "name" => "Rust",
            "version" => "",
            "mimetype" => "text/rust",
            "file_extension" => ".rs",
            // Pygments lexer, for highlighting Only needed if it differs from the 'name' field.
            // see http://pygments.org/docs/lexers/#lexers-for-the-rust-language
            "pygment_lexer" => "rust",
            // Codemirror mode, for for highlighting in the notebook. Only needed if it differs from the 'name' field.
            // codemirror use text/x-rustsrc as mimetypes
            // see https://codemirror.net/mode/rust/
            "codemirror_mode" => "rust",
        },
        "banner" => format!("EvCxR {} - Evaluation Context for Rust", env!("CARGO_PKG_VERSION")),
        "help_links" => array![
            object!{"text" => "Rust std docs",
                    "url" => "https://doc.rust-lang.org/stable/std/"}
        ],
        "status" => "ok"
    }
}
<commit_msg>allow multiple command and mixed command & code per cell in jupyter<commit_after>// Copyright 2018 Google Inc.
//
// Licensed 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
//
//     https://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.

use colored::*;
use connection::Connection;
use control_file;
use evcxr;
use evcxr::CommandContext;
use failure::Error;
use json;
use json::JsonValue;
use jupyter_message::JupyterMessage;
use std;
use std::collections::HashMap;
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use std::time;
use zmq;

// Note, to avoid potential deadlocks, each thread should lock at most one mutex at a time.
#[derive(Clone)]
pub(crate) struct Server {
    iopub: Arc<Mutex<Connection>>,
    _stdin: Arc<Mutex<Connection>>,
    latest_execution_request: Arc<Mutex<Option<JupyterMessage>>>,
    shutdown_requested_receiver: Arc<Mutex<mpsc::Receiver<()>>>,
    shutdown_requested_sender: Arc<Mutex<mpsc::Sender<()>>>,
}

impl Server {
    pub(crate) fn start(config: &control_file::Control) -> Result<Server, Error> {
        use zmq::SocketType;

        let zmq_context = zmq::Context::new();
        let heartbeat = bind_socket(config, config.hb_port, zmq_context.socket(SocketType::REP)?)?;
        let shell_socket = bind_socket(
            config,
            config.shell_port,
            zmq_context.socket(SocketType::ROUTER)?,
        )?;
        let control_socket = bind_socket(
            config,
            config.control_port,
            zmq_context.socket(SocketType::ROUTER)?,
        )?;
        let stdin_socket = bind_socket(
            config,
            config.stdin_port,
            zmq_context.socket(SocketType::ROUTER)?,
        )?;
        let iopub = Arc::new(Mutex::new(bind_socket(
            config,
            config.iopub_port,
            zmq_context.socket(SocketType::PUB)?,
        )?));

        let (shutdown_requested_sender, shutdown_requested_receiver) = mpsc::channel();

        let server = Server {
            iopub,
            latest_execution_request: Arc::new(Mutex::new(None)),
            _stdin: Arc::new(Mutex::new(stdin_socket)),
            shutdown_requested_receiver: Arc::new(Mutex::new(shutdown_requested_receiver)),
            shutdown_requested_sender: Arc::new(Mutex::new(shutdown_requested_sender)),
        };

        let (execution_sender, execution_receiver) = mpsc::channel();
        let (execution_response_sender, execution_response_receiver) = mpsc::channel();

        thread::spawn(move || Self::handle_hb(&heartbeat));
        server.start_thread(move |server: Server| server.handle_control(control_socket));
        server.start_thread(move |server: Server| {
            server.handle_shell(
                shell_socket,
                &execution_sender,
                &execution_response_receiver,
            )
        });
        let (context, outputs) = CommandContext::new()?;
        server.start_thread(move |server: Server| {
            server.handle_execution_requests(
                context,
                &execution_receiver,
                &execution_response_sender,
            )
        });
        server
            .clone()
            .start_output_pass_through_thread("stdout", outputs.stdout);
        server
            .clone()
            .start_output_pass_through_thread("stderr", outputs.stderr);
        Ok(server)
    }

    pub(crate) fn wait_for_shutdown(&self) {
        self.shutdown_requested_receiver
            .lock()
            .unwrap()
            .recv()
            .unwrap();
    }

    fn signal_shutdown(&self) {
        self.shutdown_requested_sender
            .lock()
            .unwrap()
            .send(())
            .unwrap();
    }

    fn start_thread<F>(&self, body: F)
    where
        F: FnOnce(Server) -> Result<(), Error> + std::marker::Send + 'static,
    {
        let server_clone = self.clone();
        thread::spawn(|| {
            if let Err(error) = body(server_clone) {
                eprintln!("{:?}", error);
            }
        });
    }

    fn handle_hb(connection: &Connection) -> Result<(), Error> {
        let mut message = zmq::Message::new()?;
        loop {
            connection.socket.recv(&mut message, 0)?;
            connection.socket.send(b"ping", 0)?;
        }
    }

    fn handle_execution_requests(
        self,
        mut context: CommandContext,
        receiver: &mpsc::Receiver<JupyterMessage>,
        execution_reply_sender: &mpsc::Sender<JupyterMessage>,
    ) -> Result<(), Error> {
        let mut execution_count = 1;
        loop {
            let message = receiver.recv()?;

            // If we want this clone to be cheaper, we probably only need the header, not the
            // whole message.
            *self.latest_execution_request.lock().unwrap() = Some(message.clone());
            let src = message.code();
            execution_count += 1;
            message
                .new_message("execute_input")
                .with_content(object! {
                    "execution_count" => execution_count,
                    "code" => src
                })
                .send(&mut *self.iopub.lock().unwrap())?;
            let mut has_error = false;
            for code in split_code_and_command(src) {
                // stop execution after the first error
                has_error = has_error
                    || match context.execute(&code) {
                        Ok(output) => {
                            if !output.is_empty() {
                                // Increase the odds that stdout will have been finished being sent. A
                                // less hacky alternative would be to add a print statement, then block
                                // waiting for it.
                                thread::sleep(time::Duration::from_millis(1));
                                let mut data = HashMap::new();
                                // At the time of writing the json crate appears to have a generic From
                                // implementation for a Vec<T> where T implements Into<JsonValue>. It also
                                // has conversion from HashMap<String, JsonValue>, but it doesn't have
                                // conversion from HashMap<String, T>. Perhaps send a PR? For now, we
                                // convert the values manually.
                                for (k, v) in output.content_by_mime_type {
                                    if k.contains("json") {
                                        data.insert(
                                            k,
                                            json::parse(&v).unwrap_or_else(|_| json::from(v)),
                                        );
                                    } else {
                                        data.insert(k, json::from(v));
                                    }
                                }
                                message
                                    .new_message("execute_result")
                                    .with_content(object! {
                                        "execution_count" => execution_count,
                                        "data" => data,
                                        "metadata" => HashMap::new(),
                                    })
                                    .send(&mut *self.iopub.lock().unwrap())?;
                            }
                            if let Some(duration) = output.timing {
                                // TODO replace by duration.as_millis() when stable
                                let ms =
                                    duration.as_secs() * 1000 + u64::from(duration.subsec_millis());
                                let mut data = HashMap::new();
                                data.insert(
                                    "text/html".into(),
                                    json::from(format!(
                                        "Took {}ms",
                                        ms
                                    )),
                                );
                                message
                                    .new_message("execute_result")
                                    .with_content(object! {
                                        "execution_count" => execution_count,
                                        "data" => data,
                                        "metadata" => HashMap::new(),
                                    })
                                    .send(&mut *self.iopub.lock().unwrap())?;
                            }
                            false
                        }
                        Err(errors) => {
                            self.emit_errors(&errors, &message)?;
                            true
                        }
                    };
            }
            let reply = if has_error {
                message.new_reply().with_content(object! {
                    "status" => "error",
                    "execution_count" => execution_count,
                })
            } else {
                message.new_reply().with_content(object! {
                    "status" => "ok",
                    "execution_count" => execution_count
                })
            };
            execution_reply_sender.send(reply)?;
        }
    }

    fn handle_shell(
        self,
        mut connection: Connection,
        execution_channel: &mpsc::Sender<JupyterMessage>,
        execution_reply_receiver: &mpsc::Receiver<JupyterMessage>,
    ) -> Result<(), Error> {
        loop {
            let message = JupyterMessage::read(&mut connection)?;
            // Processing of every message should be enclosed between "busy" and "idle"
            // see https://jupyter-client.readthedocs.io/en/latest/messaging.html#messages-on-the-shell-router-dealer-channel
            // Jupiter Lab doesn't use the kernel until it received "idle" for kernel_info_request
            message
                .new_message("status")
                .with_content(object! {"execution_state" => "busy"})
                .send(&mut *self.iopub.lock().unwrap())?;
            let idle = message
                .new_message("status")
                .with_content(object! {"execution_state" => "idle"});
            if message.message_type() == "kernel_info_request" {
                message
                    .new_reply()
                    .with_content(kernel_info())
                    .send(&mut connection)?;
            } else if message.message_type() == "is_complete_request" {
                message
                    .new_reply()
                    .with_content(object! {"status" => "complete"})
                    .send(&mut connection)?;
            } else if message.message_type() == "execute_request" {
                execution_channel.send(message)?;
                execution_reply_receiver.recv()?.send(&mut connection)?;
            } else if message.message_type() == "comm_open" {
                message
                    .new_message("comm_close")
                    .with_content(message.get_content().clone())
                    .send(&mut connection)?;
            } else {
                eprintln!(
                    "Got unrecognized message type on shell channel: {}",
                    message.message_type()
                );
            }
            idle.send(&mut *self.iopub.lock().unwrap())?;
        }
    }

    fn handle_control(self, mut connection: Connection) -> Result<(), Error> {
        loop {
            let message = JupyterMessage::read(&mut connection)?;
            match message.message_type() {
                "shutdown_request" => self.signal_shutdown(),
                "interrupt_request" => {
                    message.new_reply().send(&mut connection)?;
                    eprintln!(
                        "Rust doesn't support interrupting execution. Perhaps restart kernel?"
                    );
                }
                _ => {
                    eprintln!(
                        "Got unrecognized message type on control channel: {}",
                        message.message_type()
                    );
                }
            }
        }
    }

    fn start_output_pass_through_thread(
        self,
        output_name: &'static str,
        channel: mpsc::Receiver<String>,
    ) {
        thread::spawn(move || {
            while let Ok(line) = channel.recv() {
                let mut message = None;
                if let Some(exec_request) = &*self.latest_execution_request.lock().unwrap() {
                    message = Some(exec_request.new_message("stream"));
                }
                if let Some(message) = message {
                    if let Err(error) = message
                        .with_content(object! {
                            "name" => output_name,
                            "text" => format!("{}\n", line),
                        })
                        .send(&mut *self.iopub.lock().unwrap())
                    {
                        eprintln!("{}", error);
                    }
                }
            }
        });
    }

    fn emit_errors(
        &self,
        errors: &evcxr::Error,
        parent_message: &JupyterMessage,
    ) -> Result<(), Error> {
        match errors {
            evcxr::Error::CompilationErrors(errors) => {
                for error in errors {
                    let message = format!("{}", error.message().bright_red());
                    if error.is_from_user_code() {
                        let mut traceback = Vec::new();
                        for spanned_message in error.spanned_messages() {
                            for line in &spanned_message.lines {
                                traceback.push(line.clone());
                            }
                            if let Some(span) = &spanned_message.span {
                                let mut carrots = String::new();
                                for _ in 1..span.start_column {
                                    carrots.push(' ');
                                }
                                for _ in span.start_column..span.end_column {
                                    carrots.push('^');
                                }
                                traceback.push(format!(
                                    "{} {}",
                                    carrots.bright_red(),
                                    spanned_message.label.bright_blue()
                                ));
                            } else {
                                traceback.push(spanned_message.label.clone());
                            }
                        }
                        traceback.push(error.message());
                        parent_message
                            .new_message("error")
                            .with_content(object! {
                                "ename" => "Error",
                                "evalue" => error.message(),
                                "traceback" => traceback,
                            })
                            .send(&mut *self.iopub.lock().unwrap())?;
                    } else {
                        parent_message
                            .new_message("error")
                            .with_content(object! {
                                "ename" => "Error",
                                "evalue" => error.message(),
                                "traceback" => array![
                                    message
                                ],
                            })
                            .send(&mut *self.iopub.lock().unwrap())?;
                    }
                }
            }
            error => {
                let displayed_error = format!("{}", error);
                parent_message
                    .new_message("error")
                    .with_content(object! {
                        "ename" => "Error",
                        "evalue" => displayed_error.clone(),
                        "traceback" => array![displayed_error],
                    })
                    .send(&mut *self.iopub.lock().unwrap())?;
            }
        }
        Ok(())
    }
}

fn bind_socket(
    config: &control_file::Control,
    port: u16,
    socket: zmq::Socket,
) -> Result<Connection, Error> {
    let endpoint = format!("{}://{}:{}", config.transport, config.ip, port);
    socket.bind(&endpoint)?;
    Ok(Connection::new(socket, &config.key)?)
}

/// See [Kernel info documentation](https://jupyter-client.readthedocs.io/en/stable/messaging.html#kernel-info)
fn kernel_info() -> JsonValue {
    object! {
        "protocol_version" => "5.3",
        "implementation" => env!("CARGO_PKG_NAME"),
        "implementation_version" => env!("CARGO_PKG_VERSION"),
        "language_info" => object!{
            "name" => "Rust",
            "version" => "",
            "mimetype" => "text/rust",
            "file_extension" => ".rs",
            // Pygments lexer, for highlighting Only needed if it differs from the 'name' field.
            // see http://pygments.org/docs/lexers/#lexers-for-the-rust-language
            "pygment_lexer" => "rust",
            // Codemirror mode, for for highlighting in the notebook. Only needed if it differs from the 'name' field.
            // codemirror use text/x-rustsrc as mimetypes
            // see https://codemirror.net/mode/rust/
            "codemirror_mode" => "rust",
        },
        "banner" => format!("EvCxR {} - Evaluation Context for Rust", env!("CARGO_PKG_VERSION")),
        "help_links" => array![
            object!{"text" => "Rust std docs",
                    "url" => "https://doc.rust-lang.org/stable/std/"}
        ],
        "status" => "ok"
    }
}

//TODO optimize by avoiding creation of new String
fn split_code_and_command(src: &str) -> Vec<String> {
    src.lines().fold(vec![], |mut acc, l| {
        if l.starts_with(":") {
            acc.push(l.to_owned());
        } else if let Some(last) = acc.pop() {
            if !last.starts_with(":") {
                acc.push((last + "\n" + l).to_owned());
            } else {
                acc.push(last);
                acc.push(l.to_owned());
            }
        } else {
            acc.push(l.to_owned());
        }
        acc
    })
}

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

    const COMMAND_0: &str = r#":dep foo= "0.3.3""#;
    const CODE_0: &str = r#"println!(":dep code 0");"#;
    const CODE_1: &str = r#"
        println!('hello');
        eprintln!('world');
    "#;

    #[test]
    fn split_code_and_command_test_single_command() {
        let expected = vec![COMMAND_0.to_owned()];
        let actual = split_code_and_command(&expected.join("\n"));
        assert_eq!(actual, expected);
    }

    #[test]
    fn split_code_and_command_test_single_code() {
        let expected = vec![CODE_1.to_owned()];
        let actual = split_code_and_command(&expected.join("\n"));
        assert_eq!(actual, expected);
    }

    #[test]
    fn split_code_and_command_test_multi_command() {
        let expected = vec![COMMAND_0.to_owned(), COMMAND_0.to_owned()];
        let actual = split_code_and_command(&expected.join("\n"));
        assert_eq!(actual, expected);
    }

    #[test]
    fn split_code_and_command_test_mixed_command_code() {
        let expected = vec![
            COMMAND_0.to_owned(),
            COMMAND_0.to_owned(),
            CODE_0.to_owned(),
            COMMAND_0.to_owned(),
            CODE_1.to_owned(),
        ];
        let actual = split_code_and_command(&expected.join("\n"));
        assert_eq!(actual, expected);
    }
}
<|endoftext|>


<commit_before>use rustc::hir::intravisit;
use rustc::hir;
use rustc::ty;
use rustc::lint::*;
use std::collections::HashSet;
use syntax::ast;
use syntax::codemap::Span;
use utils::{span_lint, type_is_unsafe_function};

/// **What it does:** Check for functions with too many parameters.
///
/// **Why is this bad?** Functions with lots of parameters are considered bad style and reduce
/// readability (“what does the 5th parameter mean?”). Consider grouping some parameters into a
/// new type.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// fn foo(x: u32, y: u32, name: &str, c: Color, w: f32, h: f32, a: f32, b: f32) { .. }
/// ```
declare_lint! {
    pub TOO_MANY_ARGUMENTS,
    Warn,
    "functions with too many arguments"
}

/// **What it does:** Check for public functions that dereferences raw pointer arguments but are
/// not marked unsafe.
///
/// **Why is this bad?** The function should probably be marked `unsafe`, since for an arbitrary
/// raw pointer, there is no way of telling for sure if it is valid.
///
/// **Known problems:**
///
/// * It does not check functions recursively so if the pointer is passed to a private non-
/// `unsafe` function which does the dereferencing, the lint won't trigger.
/// * It only checks for arguments whose type are raw pointers, not raw pointers got from an
/// argument in some other way (`fn foo(bar: &[*const u8])` or `some_argument.get_raw_ptr()`).
///
/// **Example:**
///
/// ```rust
/// pub fn foo(x: *const u8) { println!("{}", unsafe { *x }); }
/// ```
declare_lint! {
    pub NOT_UNSAFE_PTR_ARG_DEREF,
    Warn,
    "public functions dereferencing raw pointer arguments but not marked `unsafe`"
}

#[derive(Copy,Clone)]
pub struct Functions {
    threshold: u64,
}

impl Functions {
    pub fn new(threshold: u64) -> Functions {
        Functions { threshold: threshold }
    }
}

impl LintPass for Functions {
    fn get_lints(&self) -> LintArray {
        lint_array!(TOO_MANY_ARGUMENTS, NOT_UNSAFE_PTR_ARG_DEREF)
    }
}

impl LateLintPass for Functions {
    fn check_fn(&mut self, cx: &LateContext, kind: intravisit::FnKind, decl: &hir::FnDecl, block: &hir::Block, span: Span, nodeid: ast::NodeId) {
        use rustc::hir::map::Node::*;

        let is_impl = if let Some(NodeItem(ref item)) = cx.tcx.map.find(cx.tcx.map.get_parent_node(nodeid)) {
            matches!(item.node, hir::ItemImpl(_, _, _, Some(_), _, _) | hir::ItemDefaultImpl(..))
        } else {
            false
        };

        let unsafety = match kind {
            hir::intravisit::FnKind::ItemFn(_, _, unsafety, _, _, _, _) => unsafety,
            hir::intravisit::FnKind::Method(_, sig, _, _) => sig.unsafety,
            hir::intravisit::FnKind::Closure(_) => return,
        };

        // don't warn for implementations, it's not their fault
        if !is_impl {
            self.check_arg_number(cx, decl, span);
        }

        self.check_raw_ptr(cx, unsafety, decl, block, nodeid);
    }

    fn check_trait_item(&mut self, cx: &LateContext, item: &hir::TraitItem) {
        if let hir::MethodTraitItem(ref sig, ref block) = item.node {
            self.check_arg_number(cx, &sig.decl, item.span);

            if let Some(ref block) = *block {
                self.check_raw_ptr(cx, sig.unsafety, &sig.decl, block, item.id);
            }
        }
    }
}

impl Functions {
    fn check_arg_number(&self, cx: &LateContext, decl: &hir::FnDecl, span: Span) {
        let args = decl.inputs.len() as u64;
        if args > self.threshold {
            span_lint(cx,
                      TOO_MANY_ARGUMENTS,
                      span,
                      &format!("this function has too many arguments ({}/{})", args, self.threshold));
        }
    }

    fn check_raw_ptr(&self, cx: &LateContext, unsafety: hir::Unsafety, decl: &hir::FnDecl, block: &hir::Block, nodeid: ast::NodeId) {
        if unsafety == hir::Unsafety::Normal && cx.access_levels.is_exported(nodeid) {
            let raw_ptrs = decl.inputs.iter().filter_map(|arg| raw_ptr_arg(cx, arg)).collect::<HashSet<_>>();

            if !raw_ptrs.is_empty() {
                let mut v = DerefVisitor {
                    cx: cx,
                    ptrs: raw_ptrs,
                };

                hir::intravisit::walk_block(&mut v, block);
            }
        }
    }
}

fn raw_ptr_arg(cx: &LateContext, arg: &hir::Arg) -> Option<hir::def_id::DefId> {
    if let (&hir::PatKind::Binding(_, _, _), &hir::TyPtr(_)) = (&arg.pat.node, &arg.ty.node) {
        cx.tcx.def_map.borrow().get(&arg.pat.id).map(hir::def::PathResolution::def_id)
    } else {
        None
    }
}

struct DerefVisitor<'a, 'tcx: 'a> {
    cx: &'a LateContext<'a, 'tcx>,
    ptrs: HashSet<hir::def_id::DefId>,
}

impl<'a, 'tcx, 'v> hir::intravisit::Visitor<'v> for DerefVisitor<'a, 'tcx> {
    fn visit_expr(&mut self, expr: &'v hir::Expr) {
        match expr.node {
            hir::ExprCall(ref f, ref args) => {
                let ty = self.cx.tcx.expr_ty(f);

                if type_is_unsafe_function(ty) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            }
            hir::ExprMethodCall(_, _, ref args) => {
                let method_call = ty::MethodCall::expr(expr.id);
                let base_type = self.cx.tcx.tables.borrow().method_map[&method_call].ty;

                if type_is_unsafe_function(base_type) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            }
            hir::ExprUnary(hir::UnDeref, ref ptr) => self.check_arg(ptr),
            _ => (),
        }

        hir::intravisit::walk_expr(self, expr);
    }
}

impl<'a, 'tcx: 'a> DerefVisitor<'a, 'tcx> {
    fn check_arg(&self, ptr: &hir::Expr) {
        if let Some(def) = self.cx.tcx.def_map.borrow().get(&ptr.id) {
            if self.ptrs.contains(&def.def_id()) {
                span_lint(self.cx,
                          NOT_UNSAFE_PTR_ARG_DEREF,
                          ptr.span,
                          "this public function dereferences a raw pointer but is not marked `unsafe`");
            }
        }
    }
}
<commit_msg>Make #991 work with current rust<commit_after>use rustc::hir::intravisit;
use rustc::hir;
use rustc::ty;
use rustc::lint::*;
use std::collections::HashSet;
use syntax::ast;
use syntax::codemap::Span;
use utils::{span_lint, type_is_unsafe_function};

/// **What it does:** Check for functions with too many parameters.
///
/// **Why is this bad?** Functions with lots of parameters are considered bad style and reduce
/// readability (“what does the 5th parameter mean?”). Consider grouping some parameters into a
/// new type.
///
/// **Known problems:** None.
///
/// **Example:**
///
/// ```rust
/// fn foo(x: u32, y: u32, name: &str, c: Color, w: f32, h: f32, a: f32, b: f32) { .. }
/// ```
declare_lint! {
    pub TOO_MANY_ARGUMENTS,
    Warn,
    "functions with too many arguments"
}

/// **What it does:** Check for public functions that dereferences raw pointer arguments but are
/// not marked unsafe.
///
/// **Why is this bad?** The function should probably be marked `unsafe`, since for an arbitrary
/// raw pointer, there is no way of telling for sure if it is valid.
///
/// **Known problems:**
///
/// * It does not check functions recursively so if the pointer is passed to a private non-
/// `unsafe` function which does the dereferencing, the lint won't trigger.
/// * It only checks for arguments whose type are raw pointers, not raw pointers got from an
/// argument in some other way (`fn foo(bar: &[*const u8])` or `some_argument.get_raw_ptr()`).
///
/// **Example:**
///
/// ```rust
/// pub fn foo(x: *const u8) { println!("{}", unsafe { *x }); }
/// ```
declare_lint! {
    pub NOT_UNSAFE_PTR_ARG_DEREF,
    Warn,
    "public functions dereferencing raw pointer arguments but not marked `unsafe`"
}

#[derive(Copy,Clone)]
pub struct Functions {
    threshold: u64,
}

impl Functions {
    pub fn new(threshold: u64) -> Functions {
        Functions { threshold: threshold }
    }
}

impl LintPass for Functions {
    fn get_lints(&self) -> LintArray {
        lint_array!(TOO_MANY_ARGUMENTS, NOT_UNSAFE_PTR_ARG_DEREF)
    }
}

impl LateLintPass for Functions {
    fn check_fn(&mut self, cx: &LateContext, kind: intravisit::FnKind, decl: &hir::FnDecl, block: &hir::Block, span: Span, nodeid: ast::NodeId) {
        use rustc::hir::map::Node::*;

        let is_impl = if let Some(NodeItem(ref item)) = cx.tcx.map.find(cx.tcx.map.get_parent_node(nodeid)) {
            matches!(item.node, hir::ItemImpl(_, _, _, Some(_), _, _) | hir::ItemDefaultImpl(..))
        } else {
            false
        };

        let unsafety = match kind {
            hir::intravisit::FnKind::ItemFn(_, _, unsafety, _, _, _, _) => unsafety,
            hir::intravisit::FnKind::Method(_, sig, _, _) => sig.unsafety,
            hir::intravisit::FnKind::Closure(_) => return,
        };

        // don't warn for implementations, it's not their fault
        if !is_impl {
            self.check_arg_number(cx, decl, span);
        }

        self.check_raw_ptr(cx, unsafety, decl, block, nodeid);
    }

    fn check_trait_item(&mut self, cx: &LateContext, item: &hir::TraitItem) {
        if let hir::MethodTraitItem(ref sig, ref block) = item.node {
            self.check_arg_number(cx, &sig.decl, item.span);

            if let Some(ref block) = *block {
                self.check_raw_ptr(cx, sig.unsafety, &sig.decl, block, item.id);
            }
        }
    }
}

impl Functions {
    fn check_arg_number(&self, cx: &LateContext, decl: &hir::FnDecl, span: Span) {
        let args = decl.inputs.len() as u64;
        if args > self.threshold {
            span_lint(cx,
                      TOO_MANY_ARGUMENTS,
                      span,
                      &format!("this function has too many arguments ({}/{})", args, self.threshold));
        }
    }

    fn check_raw_ptr(&self, cx: &LateContext, unsafety: hir::Unsafety, decl: &hir::FnDecl, block: &hir::Block, nodeid: ast::NodeId) {
        if unsafety == hir::Unsafety::Normal && cx.access_levels.is_exported(nodeid) {
            let raw_ptrs = decl.inputs.iter().filter_map(|arg| raw_ptr_arg(cx, arg)).collect::<HashSet<_>>();

            if !raw_ptrs.is_empty() {
                let mut v = DerefVisitor {
                    cx: cx,
                    ptrs: raw_ptrs,
                };

                hir::intravisit::walk_block(&mut v, block);
            }
        }
    }
}

fn raw_ptr_arg(cx: &LateContext, arg: &hir::Arg) -> Option<hir::def_id::DefId> {
    if let (&hir::PatKind::Binding(_, _, _), &hir::TyPtr(_)) = (&arg.pat.node, &arg.ty.node) {
        cx.tcx.def_map.borrow().get(&arg.pat.id).map(|pr| pr.full_def().def_id())
    } else {
        None
    }
}

struct DerefVisitor<'a, 'tcx: 'a> {
    cx: &'a LateContext<'a, 'tcx>,
    ptrs: HashSet<hir::def_id::DefId>,
}

impl<'a, 'tcx, 'v> hir::intravisit::Visitor<'v> for DerefVisitor<'a, 'tcx> {
    fn visit_expr(&mut self, expr: &'v hir::Expr) {
        match expr.node {
            hir::ExprCall(ref f, ref args) => {
                let ty = self.cx.tcx.expr_ty(f);

                if type_is_unsafe_function(ty) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            }
            hir::ExprMethodCall(_, _, ref args) => {
                let method_call = ty::MethodCall::expr(expr.id);
                let base_type = self.cx.tcx.tables.borrow().method_map[&method_call].ty;

                if type_is_unsafe_function(base_type) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            }
            hir::ExprUnary(hir::UnDeref, ref ptr) => self.check_arg(ptr),
            _ => (),
        }

        hir::intravisit::walk_expr(self, expr);
    }
}

impl<'a, 'tcx: 'a> DerefVisitor<'a, 'tcx> {
    fn check_arg(&self, ptr: &hir::Expr) {
        if let Some(def) = self.cx.tcx.def_map.borrow().get(&ptr.id) {
            if self.ptrs.contains(&def.full_def().def_id()) {
                span_lint(self.cx,
                          NOT_UNSAFE_PTR_ARG_DEREF,
                          ptr.span,
                          "this public function dereferences a raw pointer but is not marked `unsafe`");
            }
        }
    }
}
<|endoftext|>


<commit_before>use std;
use base::*;
use super::{LoudsTrie, NodeID, LoudsPos, LinkID, INVALID_LINK_ID};

#[derive(Copy, Clone)]
struct State<'a> {
    trie_: &'a LoudsTrie,
    node_id_: NodeID,
    louds_pos_: LoudsPos,
    link_id_: LinkID,
    key_pos_: u32,
    //key_id_: u32,
}

impl<'a> std::fmt::Debug for State<'a> {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        fmt.debug_struct("State")
            .field("node_id_", &self.node_id_)
            .field("louds_pos_", &self.louds_pos_)
            .field("link_id_", &self.link_id_)
            .field("key_pos_", &self.key_pos_)
            .finish()
    }
}

impl<'a> State<'a> {
    fn new(trie: &'a LoudsTrie, node_id: NodeID, louds_pos: LoudsPos,
           link_id: LinkID, key_pos: u32) -> State<'a> {
        State { trie_: trie, node_id_: node_id, louds_pos_: louds_pos,
                link_id_: link_id, key_pos_: key_pos }
    }
}

#[derive(Debug)]
pub struct Nav<'a> {
    history_: Vec<State<'a> >,
    key_buf_: Vec<u8>,
}

// For lookups, marisa does caching based on the input character.
// We can't do that here. May want to remove or rethink the cache
// implementation in light of this.

impl<'a> Nav<'a> {
    fn new(trie: &'a LoudsTrie) -> Nav<'a> {
        let mut out = Nav { history_: Vec::new(), key_buf_: Vec::new() };
        out.history_.push(State::new(trie, NodeID(0), LoudsPos(0),
                          INVALID_LINK_ID, 0));
        out
    }

    fn get_link_id(&self) -> LinkID {
        self.history_.last().unwrap().link_id_
    }

    fn push_str<'b>(&mut self, key: &'b[u8], trie: &'a LoudsTrie, node_id: NodeID,
                    louds_pos: LoudsPos, link_id: LinkID) {

        debug!("push_str(key: {:?})", key);

//        debug!("push_str(key: {:#?}, node_id: {:#?}, louds_pos: {:#?}, \
//                link_id: {:#?}", key, node_id, louds_pos, link_id);

        let old_len = self.key_buf_.len();
        self.key_buf_.extend(key);
        assert!(old_len <= std::u32::MAX as usize);
        self.history_.push(State::new(trie, node_id, louds_pos, link_id,
                                      old_len as u32));
    }

    fn push(&mut self, mut node_id: NodeID, louds_pos: LoudsPos) {
//        debug!("push");
        let mut trie = self.history_.last().unwrap().trie_;
        loop {
            if trie.link_flags_.at(node_id.0 as usize) {
                debug!("Link flags TRUE");
                let (linked_node_id, link_id) = trie.get_linked_ids(node_id.0
                                                                    as usize);
                // Proceed either to next trie or tail
                match &trie.next_trie_ {
                    &Some(ref next_trie) => {
                        panic!();
                        trie = &**next_trie;
                        node_id = linked_node_id; // not sure about this
                        continue;
                    },
                    &None => {
                        // FIXME: Shouldn't need this temporary vector.
                        //        'restore' should return an iterator, and
                        //        state.push should consume it.
                        let mut v = Vec::new();
                        trie.tail_.restore(linked_node_id.0 as usize, &mut v);

                        // Not sure if these values are correct/useful.
                        // If some stuff is only needed for some nodes...
                        // should reflect that in the types we use
                        self.push_str(&v, trie, node_id, louds_pos, link_id);
                    }
                }
            } else {
                debug!("Link flags FALSE");
                let node_char = [ trie.bases_[node_id.0 as usize] ];
                self.push_str(&node_char, trie, node_id, louds_pos,
                              INVALID_LINK_ID);
            }
            break;
        }
    }
    pub fn has_child(&self) -> bool {
        self.history_.last().map(|s| s.trie_.has_child(s.node_id_))
            .unwrap_or(false)
    }
    pub fn go_to_child(&mut self) -> bool {
        debug!("go_to_child");
        if let Some((node_id, louds_pos)) =
            self.history_.last()
            .and_then(|s| { s.trie_.child_pos(s.node_id_) })
        {
            debug!("  (node_id: {:?} louds_pos: {:?})", node_id.0, louds_pos.0);
            self.push(node_id, louds_pos);
            true
        }
        else {
            debug!("  no child");
            false
        }
    }
    pub fn has_prev_sibling(&self) -> bool {
        // FIXME: Is this all...?
        self.history_.last().map(|h| {
            h.trie_.louds_.at(h.louds_pos_.0 as usize - 1)
        }).unwrap_or(false)
    }
    pub fn go_to_prev_sibling(&mut self) -> bool {
        // pop history and string
        // decrease louds_pos and node_id by 1 if we have a sibling
        // get string (factor this out?)
        // push new history & string

        panic!("not implemented")
    }
    pub fn has_sibling(&self) -> bool {
        self.history_.last().map(|h| {
            h.trie_.louds_.at(h.louds_pos_.0 as usize + 1)
        }).unwrap_or(false)
    }
    pub fn go_to_sibling(&mut self) -> bool {
        debug!("go_to_sibling");
        //if self.history_.len() < 2 {
        //    return false;
        //}
        if let Some(&s) = self.history_.last() {
            //debug!("louds_pos: {:?}", s.louds_pos_.0);
            //debug!("louds_pos + 1: {:?}", s.louds_pos_.0 + 1);

            let cur_len = self.key_buf_.len();
            debug!("s.key_pos_: {:?}, cur_len: {:?}", s.key_pos_, cur_len);
            assert!((s.key_pos_ as usize) <= cur_len);
            self.key_buf_.truncate(s.key_pos_ as usize);
            if s.trie_.louds_.at(s.louds_pos_.0 as usize + 1) {
                debug!("  (node_id: {:?} louds_pos: {:?})",
                       s.node_id_.0 + 1, s.louds_pos_.0 + 1);
                self.history_.pop();
                // FIXME: What about LinkID?
                self.push(NodeID(s.node_id_.0 + 1),
                          LoudsPos(s.louds_pos_.0 + 1));
                true
            } else {
                debug!("  no sibling");
                false
            }
        } else {
            false
        }
    }
    pub fn has_parent(&self) -> bool {
        panic!("not implemented")
    }
    pub fn go_to_parent(&mut self) -> bool {
        debug!("go_to_parent");
        // Could use LOUDS-trie select1(rank0(m)) to navigate upward (within a
        // single trie), but it's probably more efficient just to keep a stack
        // and pop to go up

        //if self.history_.len() == 1 {
        //    self.history_.pop();
        //    return false;
        //}

        if let Some(s) = self.history_.pop() {
            let cur_len = self.key_buf_.len();
            assert!((s.key_pos_ as usize) <= cur_len);
            self.key_buf_.truncate(s.key_pos_ as usize);
            if let Some(s) = self.history_.last() {
                let node_id = s.node_id_;
                let louds_pos = s.louds_pos_;
                debug!("  (node_id: {:?} louds_pos: {:?})", node_id, louds_pos);
            }
            true
        } else {
            false
        }
    }
    pub fn is_leaf(&self) -> bool {
        //debug!("is_leaf");
        self.history_.last().map(|s|
            s.trie_.terminal_flags_.at(s.node_id_.0 as usize)
        ).unwrap_or(false)
    }
    //pub fn get_string(&self) -> &str {
    //    panic!("not implemented")
    //}
    pub fn get_u8(&self) -> &[u8] {
        &self.key_buf_[..]
    }
    pub fn is_end(&self) -> bool {
        panic!("not implemented")
    }
}

#[derive(Copy, Clone, Debug)]
enum DFT {
    ToChild,
    ToSibling,
    ToParentSibling,
    End
}

impl DFT {
    fn new() -> DFT {
        DFT::ToChild
    }
    fn depth_first_traversal_step<'a>(&mut self, nav: &mut Nav<'a>) -> bool {

        match *self {
            DFT::ToChild => {
                if nav.go_to_child() {
//                    debug!("{:#?}, {:#?}", *self, *nav);
                    //if nav.history_.len() == 3 { panic!() }

                    return true;
                } else {
                    *self = DFT::ToSibling;
                    return false;
                }
            },
            DFT::ToSibling => {
                if nav.go_to_sibling() {
                    *self = DFT::ToChild;
                    return true;
                } else {
                    *self = DFT::ToParentSibling;
                    return false;
                }
            },
            DFT::ToParentSibling => {
                if !nav.go_to_parent() {
                    *self = DFT::End;
                    return false;
                }
                if nav.go_to_sibling() {
                    *self = DFT::ToChild;
                    return true;
                }
                return false;
            },
            DFT::End => {
                return false;
            }
        }
    }
    fn next_terminal<'a, 'b>(&mut self, nav: &'b mut Nav<'a>)
      -> Option<&'b[u8]> {
        loop {
            match *self {
                DFT::End => { return None; },
                _ => (),
            }
            if self.depth_first_traversal_step(nav) && nav.is_leaf() {
                return Some(nav.get_u8());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use env_logger;
    use std::cmp::Ordering;
    use quickcheck as qc;
    use config::{Config, NumTries};
    use key::Key;
    use super::{DFT, Nav};
    use super::super::LoudsTrie;

    fn nav_restore_prop(v: Vec<String>, num_tries: NumTries)
      -> qc::TestResult {
        debug!("in: {:?}", v);
        let mut vu: Vec<Vec<u8>> = Vec::new();
        for s in v.iter() {
            vu.push(From::from(s.as_bytes()));
        }
        debug!("u8: {:?}", vu);

        if v.iter().any(|x| x.is_empty()) {
            debug!("");
            return qc::TestResult::discard();
        }
        let mut keys: Vec<Key> = v.iter().map(|s| Key::new(s.as_bytes()))
                                 .collect();
        let config = Config::new().with_num_tries(num_tries);
        let trie = LoudsTrie::build(&mut keys, &config);
        //debug!("trie: {:#?}", trie);

        let mut nav = Nav::new(&trie);
        let mut dft = DFT::new();

        let mut vv1: Vec<Vec<u8>> = v.iter().map(|s| From::from(s.as_bytes()))
                                    .collect();

        let mut vv2: Vec<Vec<u8>> = Vec::new();
        while let Some(s) = dft.next_terminal(&mut nav) {
            vv2.push(From::from(s));
        }

        vv1.sort();
        vv2.sort();
        debug!("vv1: {:?}", vv1);
        debug!("vv2: {:?}", vv2);

        let b = vv1.cmp(&vv2) == Ordering::Equal;
        debug!("equal: {:?}", b);
        debug!("");
        qc::TestResult::from_bool(b)
    }

    fn nav_restore_prop_1(v: Vec<String>) -> qc::TestResult {
        nav_restore_prop(v, NumTries::new(1))
    }

    #[test]
    fn nav_restore_qc() {
        let _ = env_logger::init();
        qc::quickcheck(nav_restore_prop as fn(Vec<String>, NumTries)
                       -> qc::TestResult);
    }

    #[test]
    fn nav_restore_qc_1() {
        let _ = env_logger::init();
        qc::quickcheck(nav_restore_prop_1 as fn(Vec<String>)
                       -> qc::TestResult);
    }

    fn assert_passed(tr: qc::TestResult) {
        assert!(!tr.is_failure());
    }

    #[test]
    fn nav_restore_manual() {
        let _ = env_logger::init();
        //assert_passed(nav_restore_prop_1(vec!["a".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["ab".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["\u{194}\u{128}".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["Testing".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["\u{80}".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["\u{7f}".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["~".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["\u{0}".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["Testing".to_owned(),
        //                                      "T".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["Testing".to_owned(),
        //                                      "Test".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["trouble".to_owned(),
        //                                      "Threep".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["Threep".to_owned(),
        //                                      "Test".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["trouble".to_owned(),
                                              "Threep".to_owned(),
                                              "Test".to_owned()]));
        //assert_passed(nav_restore_prop_1(vec!["Testing".to_owned(),
        //                                      "trouble".to_owned(),
        //                                      "Trouble".to_owned(),
        //                                      "Threep".to_owned(),
        //                                      "Test".to_owned()]));
    }
}

<commit_msg>restore manual test cases<commit_after>use std;
use base::*;
use super::{LoudsTrie, NodeID, LoudsPos, LinkID, INVALID_LINK_ID};

#[derive(Copy, Clone)]
struct State<'a> {
    trie_: &'a LoudsTrie,
    node_id_: NodeID,
    louds_pos_: LoudsPos,
    link_id_: LinkID,
    key_pos_: u32,
    //key_id_: u32,
}

impl<'a> std::fmt::Debug for State<'a> {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        fmt.debug_struct("State")
            .field("node_id_", &self.node_id_)
            .field("louds_pos_", &self.louds_pos_)
            .field("link_id_", &self.link_id_)
            .field("key_pos_", &self.key_pos_)
            .finish()
    }
}

impl<'a> State<'a> {
    fn new(trie: &'a LoudsTrie, node_id: NodeID, louds_pos: LoudsPos,
           link_id: LinkID, key_pos: u32) -> State<'a> {
        State { trie_: trie, node_id_: node_id, louds_pos_: louds_pos,
                link_id_: link_id, key_pos_: key_pos }
    }
}

#[derive(Debug)]
pub struct Nav<'a> {
    history_: Vec<State<'a> >,
    key_buf_: Vec<u8>,
}

// For lookups, marisa does caching based on the input character.
// We can't do that here. May want to remove or rethink the cache
// implementation in light of this.

impl<'a> Nav<'a> {
    fn new(trie: &'a LoudsTrie) -> Nav<'a> {
        let mut out = Nav { history_: Vec::new(), key_buf_: Vec::new() };
        out.history_.push(State::new(trie, NodeID(0), LoudsPos(0),
                          INVALID_LINK_ID, 0));
        out
    }

    fn get_link_id(&self) -> LinkID {
        self.history_.last().unwrap().link_id_
    }

    fn push_str<'b>(&mut self, key: &'b[u8], trie: &'a LoudsTrie, node_id: NodeID,
                    louds_pos: LoudsPos, link_id: LinkID) {

        debug!("push_str(key: {:?})", key);

//        debug!("push_str(key: {:#?}, node_id: {:#?}, louds_pos: {:#?}, \
//                link_id: {:#?}", key, node_id, louds_pos, link_id);

        let old_len = self.key_buf_.len();
        self.key_buf_.extend(key);
        assert!(old_len <= std::u32::MAX as usize);
        self.history_.push(State::new(trie, node_id, louds_pos, link_id,
                                      old_len as u32));
    }

    fn push(&mut self, mut node_id: NodeID, louds_pos: LoudsPos) {
//        debug!("push");
        let mut trie = self.history_.last().unwrap().trie_;
        loop {
            if trie.link_flags_.at(node_id.0 as usize) {
                debug!("Link flags TRUE");
                let (linked_node_id, link_id) = trie.get_linked_ids(node_id.0
                                                                    as usize);
                // Proceed either to next trie or tail
                match &trie.next_trie_ {
                    &Some(ref next_trie) => {
                        panic!();
                        trie = &**next_trie;
                        node_id = linked_node_id; // not sure about this
                        continue;
                    },
                    &None => {
                        // FIXME: Shouldn't need this temporary vector.
                        //        'restore' should return an iterator, and
                        //        state.push should consume it.
                        let mut v = Vec::new();
                        trie.tail_.restore(linked_node_id.0 as usize, &mut v);

                        // Not sure if these values are correct/useful.
                        // If some stuff is only needed for some nodes...
                        // should reflect that in the types we use
                        self.push_str(&v, trie, node_id, louds_pos, link_id);
                    }
                }
            } else {
                debug!("Link flags FALSE");
                let node_char = [ trie.bases_[node_id.0 as usize] ];
                self.push_str(&node_char, trie, node_id, louds_pos,
                              INVALID_LINK_ID);
            }
            break;
        }
    }
    pub fn has_child(&self) -> bool {
        self.history_.last().map(|s| s.trie_.has_child(s.node_id_))
            .unwrap_or(false)
    }
    pub fn go_to_child(&mut self) -> bool {
        debug!("go_to_child");
        if let Some((node_id, louds_pos)) =
            self.history_.last()
            .and_then(|s| { s.trie_.child_pos(s.node_id_) })
        {
            debug!("  (node_id: {:?} louds_pos: {:?})", node_id.0, louds_pos.0);
            self.push(node_id, louds_pos);
            true
        }
        else {
            debug!("  no child");
            false
        }
    }
    pub fn has_prev_sibling(&self) -> bool {
        // FIXME: Is this all...?
        self.history_.last().map(|h| {
            h.trie_.louds_.at(h.louds_pos_.0 as usize - 1)
        }).unwrap_or(false)
    }
    pub fn go_to_prev_sibling(&mut self) -> bool {
        // pop history and string
        // decrease louds_pos and node_id by 1 if we have a sibling
        // get string (factor this out?)
        // push new history & string

        panic!("not implemented")
    }
    pub fn has_sibling(&self) -> bool {
        self.history_.last().map(|h| {
            h.trie_.louds_.at(h.louds_pos_.0 as usize + 1)
        }).unwrap_or(false)
    }
    pub fn go_to_sibling(&mut self) -> bool {
        debug!("go_to_sibling");
        //if self.history_.len() < 2 {
        //    return false;
        //}
        if let Some(&s) = self.history_.last() {
            //debug!("louds_pos: {:?}", s.louds_pos_.0);
            //debug!("louds_pos + 1: {:?}", s.louds_pos_.0 + 1);

            let cur_len = self.key_buf_.len();
            debug!("s.key_pos_: {:?}, cur_len: {:?}", s.key_pos_, cur_len);
            assert!((s.key_pos_ as usize) <= cur_len);
            self.key_buf_.truncate(s.key_pos_ as usize);
            if s.trie_.louds_.at(s.louds_pos_.0 as usize + 1) {
                debug!("  (node_id: {:?} louds_pos: {:?})",
                       s.node_id_.0 + 1, s.louds_pos_.0 + 1);
                self.history_.pop();
                // FIXME: What about LinkID?
                self.push(NodeID(s.node_id_.0 + 1),
                          LoudsPos(s.louds_pos_.0 + 1));
                true
            } else {
                debug!("  no sibling");
                false
            }
        } else {
            false
        }
    }
    pub fn has_parent(&self) -> bool {
        panic!("not implemented")
    }
    pub fn go_to_parent(&mut self) -> bool {
        debug!("go_to_parent");
        // Could use LOUDS-trie select1(rank0(m)) to navigate upward (within a
        // single trie), but it's probably more efficient just to keep a stack
        // and pop to go up

        //if self.history_.len() == 1 {
        //    self.history_.pop();
        //    return false;
        //}

        if let Some(s) = self.history_.pop() {
            let cur_len = self.key_buf_.len();
            assert!((s.key_pos_ as usize) <= cur_len);
            self.key_buf_.truncate(s.key_pos_ as usize);
            if let Some(s) = self.history_.last() {
                let node_id = s.node_id_;
                let louds_pos = s.louds_pos_;
                debug!("  (node_id: {:?} louds_pos: {:?})", node_id, louds_pos);
            }
            true
        } else {
            false
        }
    }
    pub fn is_leaf(&self) -> bool {
        //debug!("is_leaf");
        self.history_.last().map(|s|
            s.trie_.terminal_flags_.at(s.node_id_.0 as usize)
        ).unwrap_or(false)
    }
    //pub fn get_string(&self) -> &str {
    //    panic!("not implemented")
    //}
    pub fn get_u8(&self) -> &[u8] {
        &self.key_buf_[..]
    }
    pub fn is_end(&self) -> bool {
        panic!("not implemented")
    }
}

#[derive(Copy, Clone, Debug)]
enum DFT {
    ToChild,
    ToSibling,
    ToParentSibling,
    End
}

impl DFT {
    fn new() -> DFT {
        DFT::ToChild
    }
    fn depth_first_traversal_step<'a>(&mut self, nav: &mut Nav<'a>) -> bool {

        match *self {
            DFT::ToChild => {
                if nav.go_to_child() {
//                    debug!("{:#?}, {:#?}", *self, *nav);
                    //if nav.history_.len() == 3 { panic!() }

                    return true;
                } else {
                    *self = DFT::ToSibling;
                    return false;
                }
            },
            DFT::ToSibling => {
                if nav.go_to_sibling() {
                    *self = DFT::ToChild;
                    return true;
                } else {
                    *self = DFT::ToParentSibling;
                    return false;
                }
            },
            DFT::ToParentSibling => {
                if !nav.go_to_parent() {
                    *self = DFT::End;
                    return false;
                }
                if nav.go_to_sibling() {
                    *self = DFT::ToChild;
                    return true;
                }
                return false;
            },
            DFT::End => {
                return false;
            }
        }
    }
    fn next_terminal<'a, 'b>(&mut self, nav: &'b mut Nav<'a>)
      -> Option<&'b[u8]> {
        loop {
            match *self {
                DFT::End => { return None; },
                _ => (),
            }
            if self.depth_first_traversal_step(nav) && nav.is_leaf() {
                return Some(nav.get_u8());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use env_logger;
    use std::cmp::Ordering;
    use quickcheck as qc;
    use config::{Config, NumTries};
    use key::Key;
    use super::{DFT, Nav};
    use super::super::LoudsTrie;

    fn nav_restore_prop(v: Vec<String>, num_tries: NumTries)
      -> qc::TestResult {
        debug!("in: {:?}", v);
        let mut vu: Vec<Vec<u8>> = Vec::new();
        for s in v.iter() {
            vu.push(From::from(s.as_bytes()));
        }
        debug!("u8: {:?}", vu);

        if v.iter().any(|x| x.is_empty()) {
            debug!("");
            return qc::TestResult::discard();
        }
        let mut keys: Vec<Key> = v.iter().map(|s| Key::new(s.as_bytes()))
                                 .collect();
        let config = Config::new().with_num_tries(num_tries);
        let trie = LoudsTrie::build(&mut keys, &config);
        //debug!("trie: {:#?}", trie);

        let mut nav = Nav::new(&trie);
        let mut dft = DFT::new();

        let mut vv1: Vec<Vec<u8>> = v.iter().map(|s| From::from(s.as_bytes()))
                                    .collect();

        let mut vv2: Vec<Vec<u8>> = Vec::new();
        while let Some(s) = dft.next_terminal(&mut nav) {
            vv2.push(From::from(s));
        }

        vv1.sort();
        vv2.sort();
        debug!("vv1: {:?}", vv1);
        debug!("vv2: {:?}", vv2);

        let b = vv1.cmp(&vv2) == Ordering::Equal;
        debug!("equal: {:?}", b);
        debug!("");
        qc::TestResult::from_bool(b)
    }

    fn nav_restore_prop_1(v: Vec<String>) -> qc::TestResult {
        nav_restore_prop(v, NumTries::new(1))
    }

    #[test]
    fn nav_restore_qc() {
        let _ = env_logger::init();
        qc::quickcheck(nav_restore_prop as fn(Vec<String>, NumTries)
                       -> qc::TestResult);
    }

    #[test]
    fn nav_restore_qc_1() {
        let _ = env_logger::init();
        qc::quickcheck(nav_restore_prop_1 as fn(Vec<String>)
                       -> qc::TestResult);
    }

    fn assert_passed(tr: qc::TestResult) {
        assert!(!tr.is_failure());
    }

    #[test]
    fn nav_restore_manual() {
        let _ = env_logger::init();
        assert_passed(nav_restore_prop_1(vec!["a".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["ab".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["\u{194}\u{128}".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["Testing".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["\u{80}".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["\u{7f}".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["~".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["\u{0}".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["Testing".to_owned(),
                                              "T".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["Testing".to_owned(),
                                              "Test".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["trouble".to_owned(),
                                              "Threep".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["Threep".to_owned(),
                                              "Test".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["trouble".to_owned(),
                                              "Threep".to_owned(),
                                              "Test".to_owned()]));
        assert_passed(nav_restore_prop_1(vec!["Testing".to_owned(),
                                              "trouble".to_owned(),
                                              "Trouble".to_owned(),
                                              "Threep".to_owned(),
                                              "Test".to_owned()]));
    }
}

<|endoftext|>