<commit_before>//! Postfix completion for `format`-like strings.
//!
//! `"Result {result} is {2 + 2}"` is expanded to the `"Result {} is {}", result, 2 + 2`.
//!
//! The following postfix snippets are available:
//!
//! - `format` -> `format!(...)`
//! - `println` -> `println!(...)`
//! - `log`:
//!   + `logd` -> `log::debug!(...)`
//!   + `logt` -> `log::trace!(...)`
//!   + `logi` -> `log::info!(...)`
//!   + `logw` -> `log::warn!(...)`
//!   + `loge` -> `log::error!(...)`

use super::postfix_snippet;
use crate::completion::{
    completion_config::SnippetCap, completion_context::CompletionContext,
    completion_item::Completions,
};
use syntax::ast;

pub(super) fn add_format_like_completions(
    acc: &mut Completions,
    ctx: &CompletionContext,
    dot_receiver: &ast::Expr,
    cap: SnippetCap,
    receiver_text: &str,
) {
    if !is_string_literal(receiver_text) {
        // It's not a string literal, do not parse input.
        return;
    }

    let input = &receiver_text[1..receiver_text.len() - 1];

    let mut parser = FormatStrParser::new(input);

    if parser.parse().is_ok() {
        for kind in PostfixKind::all_suggestions() {
            let snippet = parser.into_suggestion(*kind);
            let (label, detail) = kind.into_description();

            postfix_snippet(ctx, cap, &dot_receiver, label, detail, &snippet).add_to(acc);
        }
    }
}

/// Checks whether provided item is a string literal.
fn is_string_literal(item: &str) -> bool {
    if item.len() < 2 {
        return false;
    }
    if item.chars().nth(0) != Some('"') || item.chars().nth(item.len() - 1) != Some('"') {
        return false;
    }

    true
}

/// Parser for a format-like string. It is more allowing in terms of string contents,
/// as we expect variable placeholders to be filled with expressions.
#[derive(Debug)]
pub struct FormatStrParser {
    input: String,
    output: String,
    extracted_expressions: Vec<String>,
    state: State,
    parsed: bool,
}

#[derive(Debug, Clone, Copy)]
pub enum PostfixKind {
    Format,
    Panic,
    Println,
    LogDebug,
    LogTrace,
    LogInfo,
    LogWarn,
    LogError,
}

impl PostfixKind {
    pub fn all_suggestions() -> &'static [PostfixKind] {
        &[
            Self::Format,
            Self::Panic,
            Self::Println,
            Self::LogDebug,
            Self::LogTrace,
            Self::LogInfo,
            Self::LogWarn,
            Self::LogError,
        ]
    }

    pub fn into_description(self) -> (&'static str, &'static str) {
        match self {
            Self::Format => ("fmt", "format!"),
            Self::Panic => ("panic", "panic!"),
            Self::Println => ("println", "println!"),
            Self::LogDebug => ("logd", "log::debug!"),
            Self::LogTrace => ("logt", "log::trace!"),
            Self::LogInfo => ("logi", "log::info!"),
            Self::LogWarn => ("logw", "log::warn!"),
            Self::LogError => ("loge", "log::error!"),
        }
    }

    pub fn into_macro_name(self) -> &'static str {
        match self {
            Self::Format => "format!",
            Self::Panic => "panic!",
            Self::Println => "println!",
            Self::LogDebug => "log::debug!",
            Self::LogTrace => "log::trace!",
            Self::LogInfo => "log::info!",
            Self::LogWarn => "log::warn!",
            Self::LogError => "log::error!",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    NotExpr,
    MaybeExpr,
    Expr,
    MaybeIncorrect,
    FormatOpts,
}

impl FormatStrParser {
    pub fn new(input: impl Into<String>) -> Self {
        Self {
            input: input.into(),
            output: String::new(),
            extracted_expressions: Vec::new(),
            state: State::NotExpr,
            parsed: false,
        }
    }

    pub fn parse(&mut self) -> Result<(), ()> {
        let mut current_expr = String::new();

        let mut placeholder_id = 1;

        // Count of open braces inside of an expression.
        // We assume that user knows what they're doing, thus we treat it like a correct pattern, e.g.
        // "{MyStruct { val_a: 0, val_b: 1 }}".
        let mut inexpr_open_count = 0;

        for chr in self.input.chars() {
            match (self.state, chr) {
                (State::NotExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::MaybeExpr;
                }
                (State::NotExpr, '}') => {
                    self.output.push(chr);
                    self.state = State::MaybeIncorrect;
                }
                (State::NotExpr, _) => {
                    self.output.push(chr);
                }
                (State::MaybeIncorrect, '}') => {
                    // It's okay, we met "}}".
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeIncorrect, _) => {
                    // Error in the string.
                    return Err(());
                }
                (State::MaybeExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, '}') => {
                    // This is an empty sequence '{}'. Replace it with placeholder.
                    self.output.push(chr);
                    self.extracted_expressions.push(format!("${}", placeholder_id));
                    placeholder_id += 1;
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, _) => {
                    current_expr.push(chr);
                    self.state = State::Expr;
                }
                (State::Expr, '}') => {
                    if inexpr_open_count == 0 {
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::NotExpr;
                    } else {
                        // We're closing one brace met before inside of the expression.
                        current_expr.push(chr);
                        inexpr_open_count -= 1;
                    }
                }
                (State::Expr, ':') => {
                    if inexpr_open_count == 0 {
                        // We're outside of braces, thus assume that it's a specifier, like "{Some(value):?}"
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::FormatOpts;
                    } else {
                        // We're inside of braced expression, assume that it's a struct field name/value delimeter.
                        current_expr.push(chr);
                    }
                }
                (State::Expr, '{') => {
                    current_expr.push(chr);
                    inexpr_open_count += 1;
                }
                (State::Expr, _) => {
                    current_expr.push(chr);
                }
                (State::FormatOpts, '}') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::FormatOpts, _) => {
                    self.output.push(chr);
                }
            }
        }

        if self.state != State::NotExpr {
            return Err(());
        }

        self.parsed = true;
        Ok(())
    }

    pub fn into_suggestion(&self, kind: PostfixKind) -> String {
        assert!(self.parsed, "Attempt to get a suggestion from not parsed expression");

        let mut output = format!(r#"{}("{}""#, kind.into_macro_name(), self.output);
        for expr in &self.extracted_expressions {
            output += ", ";
            output += expr;
        }
        output.push(')');

        output
    }
}

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

    #[test]
    fn format_str_parser() {
        let test_vector = &[
            ("no expressions", Some(("no expressions", vec![]))),
            ("{expr} is {2 + 2}", Some(("{} is {}", vec!["expr", "2 + 2"]))),
            ("{expr:?}", Some(("{:?}", vec!["expr"]))),
            ("{malformed", None),
            ("malformed}", None),
            ("{{correct", Some(("{{correct", vec![]))),
            ("correct}}", Some(("correct}}", vec![]))),
            ("{correct}}}", Some(("{}}}", vec!["correct"]))),
            ("{correct}}}}}", Some(("{}}}}}", vec!["correct"]))),
            ("{incorrect}}", None),
            ("placeholders {} {}", Some(("placeholders {} {}", vec!["$0", "$1"]))),
            ("mixed {} {2 + 2} {}", Some(("mixed {} {} {}", vec!["$0", "2 + 2", "$1"]))),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }}",
                Some(("{}", vec!["SomeStruct { val_a: 0, val_b: 1 }"])),
            ),
            ("{expr:?} is {2.32f64:.5}", Some(("{:?} is {:.5}", vec!["expr", "2.32f64"]))),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }:?}",
                Some(("{:?}", vec!["SomeStruct { val_a: 0, val_b: 1 }"])),
            ),
            ("{     2 + 2        }", Some(("{}", vec!["2 + 2"]))),
        ];

        for (input, output) in test_vector {
            let mut parser = FormatStrParser::new(*input);
            let outcome = parser.parse();

            if let Some((result_str, result_args)) = output {
                assert!(
                    outcome.is_ok(),
                    "Outcome is error for input: {}, but the expected outcome is {:?}",
                    input,
                    output
                );
                assert_eq!(parser.output, *result_str);
                assert_eq!(&parser.extracted_expressions, result_args);
            } else {
                assert!(
                    outcome.is_err(),
                    "Outcome is OK for input: {}, but the expected outcome is error",
                    input
                );
            }
        }
    }

    #[test]
    fn test_into_suggestion() {
        let test_vector = &[
            (PostfixKind::Println, "{}", r#"println!("{}", $0)"#),
            (
                PostfixKind::LogInfo,
                "{} {expr} {} {2 + 2}",
                r#"log::info!("{} {} {} {}", $0, expr, $1, 2 + 2)"#,
            ),
            (PostfixKind::Format, "{expr:?}", r#"format!("{:?}", expr)"#),
        ];

        for (kind, input, output) in test_vector {
            let mut parser = FormatStrParser::new(*input);
            parser.parse().expect("Parsing must succeed");

            assert_eq!(&parser.into_suggestion(*kind), output);
        }
    }
}
<commit_msg>Simplify is_string_literal function<commit_after>//! Postfix completion for `format`-like strings.
//!
//! `"Result {result} is {2 + 2}"` is expanded to the `"Result {} is {}", result, 2 + 2`.
//!
//! The following postfix snippets are available:
//!
//! - `format` -> `format!(...)`
//! - `println` -> `println!(...)`
//! - `log`:
//!   + `logd` -> `log::debug!(...)`
//!   + `logt` -> `log::trace!(...)`
//!   + `logi` -> `log::info!(...)`
//!   + `logw` -> `log::warn!(...)`
//!   + `loge` -> `log::error!(...)`

use super::postfix_snippet;
use crate::completion::{
    completion_config::SnippetCap, completion_context::CompletionContext,
    completion_item::Completions,
};
use syntax::ast;

pub(super) fn add_format_like_completions(
    acc: &mut Completions,
    ctx: &CompletionContext,
    dot_receiver: &ast::Expr,
    cap: SnippetCap,
    receiver_text: &str,
) {
    if !is_string_literal(receiver_text) {
        // It's not a string literal, do not parse input.
        return;
    }

    let input = &receiver_text[1..receiver_text.len() - 1];

    let mut parser = FormatStrParser::new(input);

    if parser.parse().is_ok() {
        for kind in PostfixKind::all_suggestions() {
            let snippet = parser.into_suggestion(*kind);
            let (label, detail) = kind.into_description();

            postfix_snippet(ctx, cap, &dot_receiver, label, detail, &snippet).add_to(acc);
        }
    }
}

/// Checks whether provided item is a string literal.
fn is_string_literal(item: &str) -> bool {
    if item.len() < 2 {
        return false;
    }
    item.starts_with("\"") && item.ends_with("\"")
}

/// Parser for a format-like string. It is more allowing in terms of string contents,
/// as we expect variable placeholders to be filled with expressions.
#[derive(Debug)]
pub struct FormatStrParser {
    input: String,
    output: String,
    extracted_expressions: Vec<String>,
    state: State,
    parsed: bool,
}

#[derive(Debug, Clone, Copy)]
pub enum PostfixKind {
    Format,
    Panic,
    Println,
    LogDebug,
    LogTrace,
    LogInfo,
    LogWarn,
    LogError,
}

impl PostfixKind {
    pub fn all_suggestions() -> &'static [PostfixKind] {
        &[
            Self::Format,
            Self::Panic,
            Self::Println,
            Self::LogDebug,
            Self::LogTrace,
            Self::LogInfo,
            Self::LogWarn,
            Self::LogError,
        ]
    }

    pub fn into_description(self) -> (&'static str, &'static str) {
        match self {
            Self::Format => ("fmt", "format!"),
            Self::Panic => ("panic", "panic!"),
            Self::Println => ("println", "println!"),
            Self::LogDebug => ("logd", "log::debug!"),
            Self::LogTrace => ("logt", "log::trace!"),
            Self::LogInfo => ("logi", "log::info!"),
            Self::LogWarn => ("logw", "log::warn!"),
            Self::LogError => ("loge", "log::error!"),
        }
    }

    pub fn into_macro_name(self) -> &'static str {
        match self {
            Self::Format => "format!",
            Self::Panic => "panic!",
            Self::Println => "println!",
            Self::LogDebug => "log::debug!",
            Self::LogTrace => "log::trace!",
            Self::LogInfo => "log::info!",
            Self::LogWarn => "log::warn!",
            Self::LogError => "log::error!",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    NotExpr,
    MaybeExpr,
    Expr,
    MaybeIncorrect,
    FormatOpts,
}

impl FormatStrParser {
    pub fn new(input: impl Into<String>) -> Self {
        Self {
            input: input.into(),
            output: String::new(),
            extracted_expressions: Vec::new(),
            state: State::NotExpr,
            parsed: false,
        }
    }

    pub fn parse(&mut self) -> Result<(), ()> {
        let mut current_expr = String::new();

        let mut placeholder_id = 1;

        // Count of open braces inside of an expression.
        // We assume that user knows what they're doing, thus we treat it like a correct pattern, e.g.
        // "{MyStruct { val_a: 0, val_b: 1 }}".
        let mut inexpr_open_count = 0;

        for chr in self.input.chars() {
            match (self.state, chr) {
                (State::NotExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::MaybeExpr;
                }
                (State::NotExpr, '}') => {
                    self.output.push(chr);
                    self.state = State::MaybeIncorrect;
                }
                (State::NotExpr, _) => {
                    self.output.push(chr);
                }
                (State::MaybeIncorrect, '}') => {
                    // It's okay, we met "}}".
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeIncorrect, _) => {
                    // Error in the string.
                    return Err(());
                }
                (State::MaybeExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, '}') => {
                    // This is an empty sequence '{}'. Replace it with placeholder.
                    self.output.push(chr);
                    self.extracted_expressions.push(format!("${}", placeholder_id));
                    placeholder_id += 1;
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, _) => {
                    current_expr.push(chr);
                    self.state = State::Expr;
                }
                (State::Expr, '}') => {
                    if inexpr_open_count == 0 {
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::NotExpr;
                    } else {
                        // We're closing one brace met before inside of the expression.
                        current_expr.push(chr);
                        inexpr_open_count -= 1;
                    }
                }
                (State::Expr, ':') => {
                    if inexpr_open_count == 0 {
                        // We're outside of braces, thus assume that it's a specifier, like "{Some(value):?}"
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::FormatOpts;
                    } else {
                        // We're inside of braced expression, assume that it's a struct field name/value delimeter.
                        current_expr.push(chr);
                    }
                }
                (State::Expr, '{') => {
                    current_expr.push(chr);
                    inexpr_open_count += 1;
                }
                (State::Expr, _) => {
                    current_expr.push(chr);
                }
                (State::FormatOpts, '}') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::FormatOpts, _) => {
                    self.output.push(chr);
                }
            }
        }

        if self.state != State::NotExpr {
            return Err(());
        }

        self.parsed = true;
        Ok(())
    }

    pub fn into_suggestion(&self, kind: PostfixKind) -> String {
        assert!(self.parsed, "Attempt to get a suggestion from not parsed expression");

        let mut output = format!(r#"{}("{}""#, kind.into_macro_name(), self.output);
        for expr in &self.extracted_expressions {
            output += ", ";
            output += expr;
        }
        output.push(')');

        output
    }
}

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

    #[test]
    fn format_str_parser() {
        let test_vector = &[
            ("no expressions", Some(("no expressions", vec![]))),
            ("{expr} is {2 + 2}", Some(("{} is {}", vec!["expr", "2 + 2"]))),
            ("{expr:?}", Some(("{:?}", vec!["expr"]))),
            ("{malformed", None),
            ("malformed}", None),
            ("{{correct", Some(("{{correct", vec![]))),
            ("correct}}", Some(("correct}}", vec![]))),
            ("{correct}}}", Some(("{}}}", vec!["correct"]))),
            ("{correct}}}}}", Some(("{}}}}}", vec!["correct"]))),
            ("{incorrect}}", None),
            ("placeholders {} {}", Some(("placeholders {} {}", vec!["$1", "$2"]))),
            ("mixed {} {2 + 2} {}", Some(("mixed {} {} {}", vec!["$1", "2 + 2", "$2"]))),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }}",
                Some(("{}", vec!["SomeStruct { val_a: 0, val_b: 1 }"])),
            ),
            ("{expr:?} is {2.32f64:.5}", Some(("{:?} is {:.5}", vec!["expr", "2.32f64"]))),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }:?}",
                Some(("{:?}", vec!["SomeStruct { val_a: 0, val_b: 1 }"])),
            ),
            ("{     2 + 2        }", Some(("{}", vec!["2 + 2"]))),
        ];

        for (input, output) in test_vector {
            let mut parser = FormatStrParser::new(*input);
            let outcome = parser.parse();

            if let Some((result_str, result_args)) = output {
                assert!(
                    outcome.is_ok(),
                    "Outcome is error for input: {}, but the expected outcome is {:?}",
                    input,
                    output
                );
                assert_eq!(parser.output, *result_str);
                assert_eq!(&parser.extracted_expressions, result_args);
            } else {
                assert!(
                    outcome.is_err(),
                    "Outcome is OK for input: {}, but the expected outcome is error",
                    input
                );
            }
        }
    }

    #[test]
    fn test_into_suggestion() {
        let test_vector = &[
            (PostfixKind::Println, "{}", r#"println!("{}", $1)"#),
            (
                PostfixKind::LogInfo,
                "{} {expr} {} {2 + 2}",
                r#"log::info!("{} {} {} {}", $1, expr, $2, 2 + 2)"#,
            ),
            (PostfixKind::Format, "{expr:?}", r#"format!("{:?}", expr)"#),
        ];

        for (kind, input, output) in test_vector {
            let mut parser = FormatStrParser::new(*input);
            parser.parse().expect("Parsing must succeed");

            assert_eq!(&parser.into_suggestion(*kind), output);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io::{self, Read};
use std::os::unix::io::RawFd;
use std::sync::mpsc::{channel, Receiver, TryRecvError};
use std::thread;

use base::{error, Event, PollContext, PollToken};
use data_model::{DataInit, Le16, Le32};
use vm_memory::GuestMemory;

use super::{
    base_features, copy_config, Interrupt, Queue, Reader, VirtioDevice, Writer, TYPE_CONSOLE,
};
use crate::SerialDevice;

const QUEUE_SIZE: u16 = 256;

// For now, just implement port 0 (receiveq and transmitq).
// If VIRTIO_CONSOLE_F_MULTIPORT is implemented, more queues will be needed.
const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE, QUEUE_SIZE];

#[derive(Copy, Clone, Debug, Default)]
#[repr(C)]
struct virtio_console_config {
    cols: Le16,
    rows: Le16,
    max_nr_ports: Le32,
    emerg_wr: Le32,
}

// Safe because it only has data and has no implicit padding.
unsafe impl DataInit for virtio_console_config {}

struct Worker {
    mem: GuestMemory,
    interrupt: Interrupt,
    input: Option<Box<dyn io::Read + Send>>,
    output: Option<Box<dyn io::Write + Send>>,
}

fn write_output(output: &mut Box<dyn io::Write>, data: &[u8]) -> io::Result<()> {
    output.write_all(&data)?;
    output.flush()
}

impl Worker {
    fn process_transmit_request(
        mut reader: Reader,
        output: &mut Box<dyn io::Write>,
    ) -> io::Result<u32> {
        let len = reader.available_bytes();
        let mut data = vec![0u8; len];
        reader.read_exact(&mut data)?;
        write_output(output, &data)?;
        Ok(0)
    }

    fn process_transmit_queue(
        &mut self,
        transmit_queue: &mut Queue,
        output: &mut Box<dyn io::Write>,
    ) {
        let mut needs_interrupt = false;
        while let Some(avail_desc) = transmit_queue.pop(&self.mem) {
            let desc_index = avail_desc.index;

            let reader = match Reader::new(self.mem.clone(), avail_desc) {
                Ok(r) => r,
                Err(e) => {
                    error!("console: failed to create reader: {}", e);
                    transmit_queue.add_used(&self.mem, desc_index, 0);
                    needs_interrupt = true;
                    continue;
                }
            };

            let len = match Self::process_transmit_request(reader, output) {
                Ok(written) => written,
                Err(e) => {
                    error!("console: process_transmit_request failed: {}", e);
                    0
                }
            };

            transmit_queue.add_used(&self.mem, desc_index, len);
            needs_interrupt = true;
        }

        if needs_interrupt {
            self.interrupt.signal_used_queue(transmit_queue.vector);
        }
    }

    // Start a thread that reads self.input and sends the input back via the returned channel.
    //
    // `in_avail_evt` will be triggered by the thread when new input is available.
    fn spawn_input_thread(&mut self, in_avail_evt: &Event) -> Option<Receiver<u8>> {
        let mut rx = match self.input.take() {
            Some(input) => input,
            None => return None,
        };

        let (send_channel, recv_channel) = channel();

        let thread_in_avail_evt = match in_avail_evt.try_clone() {
            Ok(evt) => evt,
            Err(e) => {
                error!("failed to clone in_avail_evt: {}", e);
                return None;
            }
        };

        // The input thread runs in detached mode and will exit when channel is disconnected because
        // the console device has been dropped.
        let res = thread::Builder::new()
            .name(format!("console_input"))
            .spawn(move || {
                let mut rx_buf = [0u8; 1];
                loop {
                    match rx.read(&mut rx_buf) {
                        Ok(0) => break, // Assume the stream of input has ended.
                        Ok(_) => {
                            if send_channel.send(rx_buf[0]).is_err() {
                                // The receiver has disconnected.
                                break;
                            }
                            thread_in_avail_evt.write(1).unwrap();
                        }
                        Err(e) => {
                            // Being interrupted is not an error, but everything else is.
                            if e.kind() != io::ErrorKind::Interrupted {
                                error!(
                                    "failed to read for bytes to queue into console device: {}",
                                    e
                                );
                                break;
                            }
                        }
                    }
                }
            });
        if let Err(e) = res {
            error!("failed to spawn input thread: {}", e);
            return None;
        }
        Some(recv_channel)
    }

    // Check for input from `in_channel_opt` and transfer it to the receive queue, if any.
    fn handle_input(
        &mut self,
        in_channel_opt: &mut Option<Receiver<u8>>,
        receive_queue: &mut Queue,
    ) {
        let in_channel = match in_channel_opt.as_ref() {
            Some(v) => v,
            None => return,
        };

        while let Some(desc) = receive_queue.peek(&self.mem) {
            let desc_index = desc.index;
            let mut writer = match Writer::new(self.mem.clone(), desc) {
                Ok(w) => w,
                Err(e) => {
                    error!("console: failed to create Writer: {}", e);
                    break;
                }
            };

            let mut disconnected = false;
            while writer.available_bytes() > 0 {
                match in_channel.try_recv() {
                    Ok(byte) => {
                        writer.write_obj(byte).unwrap();
                    }
                    Err(TryRecvError::Empty) => break,
                    Err(TryRecvError::Disconnected) => {
                        disconnected = true;
                        break;
                    }
                }
            }

            let bytes_written = writer.bytes_written() as u32;

            if bytes_written > 0 {
                receive_queue.pop_peeked(&self.mem);
                receive_queue.add_used(&self.mem, desc_index, bytes_written);
                self.interrupt.signal_used_queue(receive_queue.vector);
            }

            if disconnected {
                // Set in_channel to None so that future handle_input calls exit early.
                in_channel_opt.take();
                return;
            }

            if bytes_written == 0 {
                break;
            }
        }
    }

    fn run(&mut self, mut queues: Vec<Queue>, mut queue_evts: Vec<Event>, kill_evt: Event) {
        #[derive(PollToken)]
        enum Token {
            ReceiveQueueAvailable,
            TransmitQueueAvailable,
            InputAvailable,
            InterruptResample,
            Kill,
        }

        // Device -> driver
        let (mut receive_queue, receive_evt) = (queues.remove(0), queue_evts.remove(0));

        // Driver -> device
        let (mut transmit_queue, transmit_evt) = (queues.remove(0), queue_evts.remove(0));

        let in_avail_evt = match Event::new() {
            Ok(evt) => evt,
            Err(e) => {
                error!("failed creating Event: {}", e);
                return;
            }
        };

        // Spawn a separate thread to poll self.input.
        // A thread is used because io::Read only provides a blocking interface, and there is no
        // generic way to add an io::Read instance to a poll context (it may not be backed by a file
        // descriptor).  Moving the blocking read call to a separate thread and sending data back to
        // the main worker thread with an event for notification bridges this gap.
        let mut in_channel = self.spawn_input_thread(&in_avail_evt);

        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
            (&transmit_evt, Token::TransmitQueueAvailable),
            (&receive_evt, Token::ReceiveQueueAvailable),
            (&in_avail_evt, Token::InputAvailable),
            (self.interrupt.get_resample_evt(), Token::InterruptResample),
            (&kill_evt, Token::Kill),
        ]) {
            Ok(pc) => pc,
            Err(e) => {
                error!("failed creating PollContext: {}", e);
                return;
            }
        };

        let mut output: Box<dyn io::Write> = match self.output.take() {
            Some(o) => o,
            None => Box::new(io::sink()),
        };

        'poll: loop {
            let events = match poll_ctx.wait() {
                Ok(v) => v,
                Err(e) => {
                    error!("failed polling for events: {}", e);
                    break;
                }
            };

            for event in events.iter_readable() {
                match event.token() {
                    Token::TransmitQueueAvailable => {
                        if let Err(e) = transmit_evt.read() {
                            error!("failed reading transmit queue Event: {}", e);
                            break 'poll;
                        }
                        self.process_transmit_queue(&mut transmit_queue, &mut output);
                    }
                    Token::ReceiveQueueAvailable => {
                        if let Err(e) = receive_evt.read() {
                            error!("failed reading receive queue Event: {}", e);
                            break 'poll;
                        }
                        self.handle_input(&mut in_channel, &mut receive_queue);
                    }
                    Token::InputAvailable => {
                        if let Err(e) = in_avail_evt.read() {
                            error!("failed reading in_avail_evt: {}", e);
                            break 'poll;
                        }
                        self.handle_input(&mut in_channel, &mut receive_queue);
                    }
                    Token::InterruptResample => {
                        self.interrupt.interrupt_resample();
                    }
                    Token::Kill => break 'poll,
                }
            }
        }
    }
}

/// Virtio console device.
pub struct Console {
    base_features: u64,
    kill_evt: Option<Event>,
    worker_thread: Option<thread::JoinHandle<Worker>>,
    input: Option<Box<dyn io::Read + Send>>,
    output: Option<Box<dyn io::Write + Send>>,
    keep_fds: Vec<RawFd>,
}

impl SerialDevice for Console {
    fn new(
        _evt: Event,
        input: Option<Box<dyn io::Read + Send>>,
        output: Option<Box<dyn io::Write + Send>>,
        keep_fds: Vec<RawFd>,
    ) -> Console {
        Console {
            base_features: base_features(),
            kill_evt: None,
            worker_thread: None,
            input,
            output,
            keep_fds,
        }
    }
}

impl Drop for Console {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            // Ignore the result because there is nothing we can do about it.
            let _ = kill_evt.write(1);
        }

        if let Some(worker_thread) = self.worker_thread.take() {
            let _ = worker_thread.join();
        }
    }
}

impl VirtioDevice for Console {
    fn keep_fds(&self) -> Vec<RawFd> {
        self.keep_fds.clone()
    }

    fn features(&self) -> u64 {
        self.base_features
    }

    fn device_type(&self) -> u32 {
        TYPE_CONSOLE
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config = virtio_console_config {
            max_nr_ports: 1.into(),
            ..Default::default()
        };
        copy_config(data, 0, config.as_slice(), offset);
    }

    fn activate(
        &mut self,
        mem: GuestMemory,
        interrupt: Interrupt,
        queues: Vec<Queue>,
        queue_evts: Vec<Event>,
    ) {
        if queues.len() < 2 || queue_evts.len() < 2 {
            return;
        }

        let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e))) {
            Ok(v) => v,
            Err(e) => {
                error!("failed creating kill Event pair: {}", e);
                return;
            }
        };
        self.kill_evt = Some(self_kill_evt);

        let input = self.input.take();
        let output = self.output.take();

        let worker_result = thread::Builder::new()
            .name("virtio_console".to_string())
            .spawn(move || {
                let mut worker = Worker {
                    mem,
                    interrupt,
                    input,
                    output,
                };
                worker.run(queues, queue_evts, kill_evt);
                worker
            });

        match worker_result {
            Err(e) => {
                error!("failed to spawn virtio_console worker: {}", e);
                return;
            }
            Ok(join_handle) => {
                self.worker_thread = Some(join_handle);
            }
        }
    }

    fn reset(&mut self) -> bool {
        if let Some(kill_evt) = self.kill_evt.take() {
            if kill_evt.write(1).is_err() {
                error!("{}: failed to notify the kill event", self.debug_label());
                return false;
            }
        }

        if let Some(worker_thread) = self.worker_thread.take() {
            match worker_thread.join() {
                Err(_) => {
                    error!("{}: failed to get back resources", self.debug_label());
                    return false;
                }
                Ok(worker) => {
                    self.input = worker.input;
                    self.output = worker.output;
                    return true;
                }
            }
        }
        false
    }
}
<commit_msg>Use larger read buffers for virtio-console.<commit_after>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io::{self, Read, Write};
use std::os::unix::io::RawFd;
use std::sync::mpsc::{channel, Receiver, TryRecvError};
use std::thread;

use base::{error, Event, PollContext, PollToken};
use data_model::{DataInit, Le16, Le32};
use vm_memory::GuestMemory;

use super::{
    base_features, copy_config, Interrupt, Queue, Reader, VirtioDevice, Writer, TYPE_CONSOLE,
};
use crate::SerialDevice;

const QUEUE_SIZE: u16 = 256;

// For now, just implement port 0 (receiveq and transmitq).
// If VIRTIO_CONSOLE_F_MULTIPORT is implemented, more queues will be needed.
const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE, QUEUE_SIZE];

#[derive(Copy, Clone, Debug, Default)]
#[repr(C)]
struct virtio_console_config {
    cols: Le16,
    rows: Le16,
    max_nr_ports: Le32,
    emerg_wr: Le32,
}

// Safe because it only has data and has no implicit padding.
unsafe impl DataInit for virtio_console_config {}

struct Worker {
    mem: GuestMemory,
    interrupt: Interrupt,
    input: Option<Box<dyn io::Read + Send>>,
    output: Option<Box<dyn io::Write + Send>>,
}

fn write_output(output: &mut Box<dyn io::Write>, data: &[u8]) -> io::Result<()> {
    output.write_all(&data)?;
    output.flush()
}

impl Worker {
    fn process_transmit_request(
        mut reader: Reader,
        output: &mut Box<dyn io::Write>,
    ) -> io::Result<u32> {
        let len = reader.available_bytes();
        let mut data = vec![0u8; len];
        reader.read_exact(&mut data)?;
        write_output(output, &data)?;
        Ok(0)
    }

    fn process_transmit_queue(
        &mut self,
        transmit_queue: &mut Queue,
        output: &mut Box<dyn io::Write>,
    ) {
        let mut needs_interrupt = false;
        while let Some(avail_desc) = transmit_queue.pop(&self.mem) {
            let desc_index = avail_desc.index;

            let reader = match Reader::new(self.mem.clone(), avail_desc) {
                Ok(r) => r,
                Err(e) => {
                    error!("console: failed to create reader: {}", e);
                    transmit_queue.add_used(&self.mem, desc_index, 0);
                    needs_interrupt = true;
                    continue;
                }
            };

            let len = match Self::process_transmit_request(reader, output) {
                Ok(written) => written,
                Err(e) => {
                    error!("console: process_transmit_request failed: {}", e);
                    0
                }
            };

            transmit_queue.add_used(&self.mem, desc_index, len);
            needs_interrupt = true;
        }

        if needs_interrupt {
            self.interrupt.signal_used_queue(transmit_queue.vector);
        }
    }

    // Start a thread that reads self.input and sends the input back via the returned channel.
    //
    // `in_avail_evt` will be triggered by the thread when new input is available.
    fn spawn_input_thread(&mut self, in_avail_evt: &Event) -> Option<Receiver<Vec<u8>>> {
        let mut rx = match self.input.take() {
            Some(input) => input,
            None => return None,
        };

        let (send_channel, recv_channel) = channel();

        let thread_in_avail_evt = match in_avail_evt.try_clone() {
            Ok(evt) => evt,
            Err(e) => {
                error!("failed to clone in_avail_evt: {}", e);
                return None;
            }
        };

        // The input thread runs in detached mode and will exit when channel is disconnected because
        // the console device has been dropped.
        let res = thread::Builder::new()
            .name(format!("console_input"))
            .spawn(move || {
                loop {
                    let mut rx_buf = vec![0u8; 1 << 12];
                    match rx.read(&mut rx_buf) {
                        Ok(0) => break, // Assume the stream of input has ended.
                        Ok(size) => {
                            rx_buf.truncate(size);
                            if send_channel.send(rx_buf).is_err() {
                                // The receiver has disconnected.
                                break;
                            }
                            thread_in_avail_evt.write(1).unwrap();
                        }
                        Err(e) => {
                            // Being interrupted is not an error, but everything else is.
                            if e.kind() != io::ErrorKind::Interrupted {
                                error!(
                                    "failed to read for bytes to queue into console device: {}",
                                    e
                                );
                                break;
                            }
                        }
                    }
                }
            });
        if let Err(e) = res {
            error!("failed to spawn input thread: {}", e);
            return None;
        }
        Some(recv_channel)
    }

    // Check for input from `in_channel_opt` and transfer it to the receive queue, if any.
    fn handle_input(
        &mut self,
        in_channel_opt: &mut Option<Receiver<Vec<u8>>>,
        receive_queue: &mut Queue,
    ) {
        let in_channel = match in_channel_opt.as_ref() {
            Some(v) => v,
            None => return,
        };

        while let Some(desc) = receive_queue.peek(&self.mem) {
            let desc_index = desc.index;
            let mut writer = match Writer::new(self.mem.clone(), desc) {
                Ok(w) => w,
                Err(e) => {
                    error!("console: failed to create Writer: {}", e);
                    break;
                }
            };

            let mut disconnected = false;
            while writer.available_bytes() > 0 {
                match in_channel.try_recv() {
                    Ok(data) => {
                        writer.write_all(&data).unwrap();
                    }
                    Err(TryRecvError::Empty) => break,
                    Err(TryRecvError::Disconnected) => {
                        disconnected = true;
                        break;
                    }
                }
            }

            let bytes_written = writer.bytes_written() as u32;

            if bytes_written > 0 {
                receive_queue.pop_peeked(&self.mem);
                receive_queue.add_used(&self.mem, desc_index, bytes_written);
                self.interrupt.signal_used_queue(receive_queue.vector);
            }

            if disconnected {
                // Set in_channel to None so that future handle_input calls exit early.
                in_channel_opt.take();
                return;
            }

            if bytes_written == 0 {
                break;
            }
        }
    }

    fn run(&mut self, mut queues: Vec<Queue>, mut queue_evts: Vec<Event>, kill_evt: Event) {
        #[derive(PollToken)]
        enum Token {
            ReceiveQueueAvailable,
            TransmitQueueAvailable,
            InputAvailable,
            InterruptResample,
            Kill,
        }

        // Device -> driver
        let (mut receive_queue, receive_evt) = (queues.remove(0), queue_evts.remove(0));

        // Driver -> device
        let (mut transmit_queue, transmit_evt) = (queues.remove(0), queue_evts.remove(0));

        let in_avail_evt = match Event::new() {
            Ok(evt) => evt,
            Err(e) => {
                error!("failed creating Event: {}", e);
                return;
            }
        };

        // Spawn a separate thread to poll self.input.
        // A thread is used because io::Read only provides a blocking interface, and there is no
        // generic way to add an io::Read instance to a poll context (it may not be backed by a file
        // descriptor).  Moving the blocking read call to a separate thread and sending data back to
        // the main worker thread with an event for notification bridges this gap.
        let mut in_channel = self.spawn_input_thread(&in_avail_evt);

        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
            (&transmit_evt, Token::TransmitQueueAvailable),
            (&receive_evt, Token::ReceiveQueueAvailable),
            (&in_avail_evt, Token::InputAvailable),
            (self.interrupt.get_resample_evt(), Token::InterruptResample),
            (&kill_evt, Token::Kill),
        ]) {
            Ok(pc) => pc,
            Err(e) => {
                error!("failed creating PollContext: {}", e);
                return;
            }
        };

        let mut output: Box<dyn io::Write> = match self.output.take() {
            Some(o) => o,
            None => Box::new(io::sink()),
        };

        'poll: loop {
            let events = match poll_ctx.wait() {
                Ok(v) => v,
                Err(e) => {
                    error!("failed polling for events: {}", e);
                    break;
                }
            };

            for event in events.iter_readable() {
                match event.token() {
                    Token::TransmitQueueAvailable => {
                        if let Err(e) = transmit_evt.read() {
                            error!("failed reading transmit queue Event: {}", e);
                            break 'poll;
                        }
                        self.process_transmit_queue(&mut transmit_queue, &mut output);
                    }
                    Token::ReceiveQueueAvailable => {
                        if let Err(e) = receive_evt.read() {
                            error!("failed reading receive queue Event: {}", e);
                            break 'poll;
                        }
                        self.handle_input(&mut in_channel, &mut receive_queue);
                    }
                    Token::InputAvailable => {
                        if let Err(e) = in_avail_evt.read() {
                            error!("failed reading in_avail_evt: {}", e);
                            break 'poll;
                        }
                        self.handle_input(&mut in_channel, &mut receive_queue);
                    }
                    Token::InterruptResample => {
                        self.interrupt.interrupt_resample();
                    }
                    Token::Kill => break 'poll,
                }
            }
        }
    }
}

/// Virtio console device.
pub struct Console {
    base_features: u64,
    kill_evt: Option<Event>,
    worker_thread: Option<thread::JoinHandle<Worker>>,
    input: Option<Box<dyn io::Read + Send>>,
    output: Option<Box<dyn io::Write + Send>>,
    keep_fds: Vec<RawFd>,
}

impl SerialDevice for Console {
    fn new(
        _evt: Event,
        input: Option<Box<dyn io::Read + Send>>,
        output: Option<Box<dyn io::Write + Send>>,
        keep_fds: Vec<RawFd>,
    ) -> Console {
        Console {
            base_features: base_features(),
            kill_evt: None,
            worker_thread: None,
            input,
            output,
            keep_fds,
        }
    }
}

impl Drop for Console {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            // Ignore the result because there is nothing we can do about it.
            let _ = kill_evt.write(1);
        }

        if let Some(worker_thread) = self.worker_thread.take() {
            let _ = worker_thread.join();
        }
    }
}

impl VirtioDevice for Console {
    fn keep_fds(&self) -> Vec<RawFd> {
        self.keep_fds.clone()
    }

    fn features(&self) -> u64 {
        self.base_features
    }

    fn device_type(&self) -> u32 {
        TYPE_CONSOLE
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config = virtio_console_config {
            max_nr_ports: 1.into(),
            ..Default::default()
        };
        copy_config(data, 0, config.as_slice(), offset);
    }

    fn activate(
        &mut self,
        mem: GuestMemory,
        interrupt: Interrupt,
        queues: Vec<Queue>,
        queue_evts: Vec<Event>,
    ) {
        if queues.len() < 2 || queue_evts.len() < 2 {
            return;
        }

        let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e))) {
            Ok(v) => v,
            Err(e) => {
                error!("failed creating kill Event pair: {}", e);
                return;
            }
        };
        self.kill_evt = Some(self_kill_evt);

        let input = self.input.take();
        let output = self.output.take();

        let worker_result = thread::Builder::new()
            .name("virtio_console".to_string())
            .spawn(move || {
                let mut worker = Worker {
                    mem,
                    interrupt,
                    input,
                    output,
                };
                worker.run(queues, queue_evts, kill_evt);
                worker
            });

        match worker_result {
            Err(e) => {
                error!("failed to spawn virtio_console worker: {}", e);
                return;
            }
            Ok(join_handle) => {
                self.worker_thread = Some(join_handle);
            }
        }
    }

    fn reset(&mut self) -> bool {
        if let Some(kill_evt) = self.kill_evt.take() {
            if kill_evt.write(1).is_err() {
                error!("{}: failed to notify the kill event", self.debug_label());
                return false;
            }
        }

        if let Some(worker_thread) = self.worker_thread.take() {
            match worker_thread.join() {
                Err(_) => {
                    error!("{}: failed to get back resources", self.debug_label());
                    return false;
                }
                Ok(worker) => {
                    self.input = worker.input;
                    self.output = worker.output;
                    return true;
                }
            }
        }
        false
    }
}
<|endoftext|>


<commit_before>macro_rules! not_none {
    ($bytes:expr) => {
        match $bytes {
            Some(bytes) => bytes,
            None => return Err(Box::new($crate::types::impls::option::UnexpectedNullError {
                msg: "Unexpected null for non-null column".to_string(),
            })),
        }
    }
}

macro_rules! expression_impls {
    ($($Source:ident -> $Target:ty),+,) => {
        $(
            impl<'a> $crate::expression::AsExpression<types::$Source> for $Target {
                type Expression = $crate::expression::bound::Bound<types::$Source, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a, 'expr> $crate::expression::AsExpression<types::$Source> for &'expr $Target {
                type Expression = $crate::expression::bound::Bound<types::$Source, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a> $crate::expression::AsExpression<$crate::types::Nullable<types::$Source>> for $Target {
                type Expression = $crate::expression::bound::Bound<$crate::types::Nullable<types::$Source>, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a, 'expr> $crate::expression::AsExpression<$crate::types::Nullable<types::$Source>> for &'expr $Target {
                type Expression = $crate::expression::bound::Bound<$crate::types::Nullable<types::$Source>, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a, DB> $crate::types::ToSql<$crate::types::Nullable<types::$Source>, DB> for $Target where
                DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
                $Target: $crate::types::ToSql<types::$Source, DB>,
            {
                fn to_sql<W: ::std::io::Write>(&self, out: &mut W) -> Result<$crate::types::IsNull, Box<::std::error::Error+Send+Sync>> {
                    $crate::types::ToSql::<types::$Source, DB>::to_sql(self, out)
                }
            }
        )+
    }
}

macro_rules! queryable_impls {
    ($($Source:ident -> $Target:ty),+,) => {$(
        impl<DB> $crate::types::FromSqlRow<types::$Source, DB> for $Target where
            DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
            $Target: $crate::types::FromSql<types::$Source, DB>,
        {
            fn build_from_row<R: $crate::row::Row<DB>>(row: &mut R) -> Result<Self, Box<::std::error::Error+Send+Sync>> {
                $crate::types::FromSql::<types::$Source, DB>::from_sql(row.take())
            }
        }

        #[cfg(not(feature = "unstable"))]
        impl<DB> $crate::types::FromSqlRow<$crate::types::Nullable<types::$Source>, DB> for Option<$Target> where
            DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
            Option<$Target>: $crate::types::FromSql<$crate::types::Nullable<types::$Source>, DB>,
        {
            fn build_from_row<R: $crate::row::Row<DB>>(row: &mut R) -> Result<Self, Box<::std::error::Error+Send+Sync>> {
                $crate::types::FromSql::<$crate::types::Nullable<types::$Source>, DB>::from_sql(row.take())
            }
        }

        impl<DB> $crate::query_source::Queryable<types::$Source, DB> for $Target where
            DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
            $Target: $crate::types::FromSqlRow<types::$Source, DB>,
        {
            type Row = Self;

            fn build(row: Self::Row) -> Self {
                row
            }
        }
    )+}
}

macro_rules! primitive_impls {
    ($Source:ident -> ($Target:ty, pg: ($oid:expr, $array_oid:expr), sqlite: ($tpe:ident))) => {
        #[cfg(feature = "sqlite")]
        impl types::HasSqlType<types::$Source> for $crate::sqlite::Sqlite {
            fn metadata() -> $crate::sqlite::SqliteType {
                $crate::sqlite::SqliteType::$tpe
            }
        }

        primitive_impls!($Source -> ($Target, pg: ($oid, $array_oid)));
    };

    ($Source:ident -> ($Target:ty, pg: ($oid:expr, $array_oid:expr))) => {
        primitive_impls!($Source -> (pg: ($oid, $array_oid)));
        primitive_impls!($Source -> $Target);
    };

    ($Source:ident -> (pg: ($oid:expr, $array_oid:expr))) => {
        #[cfg(feature = "postgres")]
        impl types::HasSqlType<types::$Source> for $crate::pg::Pg {
            fn metadata() -> $crate::pg::PgTypeMetadata {
                $crate::pg::PgTypeMetadata {
                    oid: $oid,
                    array_oid: $array_oid,
                }
            }
        }
    };

    ($Source:ident -> $Target:ty) => {
        primitive_impls!($Source);
        queryable_impls!($Source -> $Target,);
        expression_impls!($Source -> $Target,);
    };

    ($Source:ident) => {
        impl types::HasSqlType<types::$Source> for $crate::backend::Debug {
            fn metadata() {}
        }

        impl $crate::query_builder::QueryId for types::$Source {
            type QueryId = Self;

            fn has_static_query_id() -> bool {
                true
            }
        }

        impl types::NotNull for types::$Source {
        }
    }
}

macro_rules! debug_to_sql {
    ($sql_type:ty, $ty:ty) => {
        impl $crate::types::ToSql<$sql_type, $crate::backend::Debug> for $ty {
            fn to_sql<W: ::std::io::Write>(&self, _: &mut W) -> Result<$crate::types::IsNull, Box<::std::error::Error+Send+Sync>> {
                Ok($crate::types::IsNull::No)
            }
        }
    };
}

mod date_and_time;
pub mod floats;
mod integers;
pub mod option;
mod primitives;
mod tuples;
<commit_msg>Make macros public to allow easily implementing traits outside of diesel<commit_after>macro_rules! not_none {
    ($bytes:expr) => {
        match $bytes {
            Some(bytes) => bytes,
            None => return Err(Box::new($crate::types::impls::option::UnexpectedNullError {
                msg: "Unexpected null for non-null column".to_string(),
            })),
        }
    }
}

#[doc(hidden)]
#[macro_export]
macro_rules! expression_impls {
    ($($Source:ident -> $Target:ty),+,) => {
        $(
            impl<'a> $crate::expression::AsExpression<types::$Source> for $Target {
                type Expression = $crate::expression::bound::Bound<types::$Source, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a, 'expr> $crate::expression::AsExpression<types::$Source> for &'expr $Target {
                type Expression = $crate::expression::bound::Bound<types::$Source, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a> $crate::expression::AsExpression<$crate::types::Nullable<types::$Source>> for $Target {
                type Expression = $crate::expression::bound::Bound<$crate::types::Nullable<types::$Source>, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a, 'expr> $crate::expression::AsExpression<$crate::types::Nullable<types::$Source>> for &'expr $Target {
                type Expression = $crate::expression::bound::Bound<$crate::types::Nullable<types::$Source>, Self>;

                fn as_expression(self) -> Self::Expression {
                    $crate::expression::bound::Bound::new(self)
                }
            }

            impl<'a, DB> $crate::types::ToSql<$crate::types::Nullable<types::$Source>, DB> for $Target where
                DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
                $Target: $crate::types::ToSql<types::$Source, DB>,
            {
                fn to_sql<W: ::std::io::Write>(&self, out: &mut W) -> Result<$crate::types::IsNull, Box<::std::error::Error+Send+Sync>> {
                    $crate::types::ToSql::<types::$Source, DB>::to_sql(self, out)
                }
            }
        )+
    }
}

#[doc(hidden)]
#[macro_export]
macro_rules! queryable_impls {
    ($($Source:ident -> $Target:ty),+,) => {$(
        impl<DB> $crate::types::FromSqlRow<types::$Source, DB> for $Target where
            DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
            $Target: $crate::types::FromSql<types::$Source, DB>,
        {
            fn build_from_row<R: $crate::row::Row<DB>>(row: &mut R) -> Result<Self, Box<::std::error::Error+Send+Sync>> {
                $crate::types::FromSql::<types::$Source, DB>::from_sql(row.take())
            }
        }

        #[cfg(not(feature = "unstable"))]
        impl<DB> $crate::types::FromSqlRow<$crate::types::Nullable<types::$Source>, DB> for Option<$Target> where
            DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
            Option<$Target>: $crate::types::FromSql<$crate::types::Nullable<types::$Source>, DB>,
        {
            fn build_from_row<R: $crate::row::Row<DB>>(row: &mut R) -> Result<Self, Box<::std::error::Error+Send+Sync>> {
                $crate::types::FromSql::<$crate::types::Nullable<types::$Source>, DB>::from_sql(row.take())
            }
        }

        impl<DB> $crate::query_source::Queryable<types::$Source, DB> for $Target where
            DB: $crate::backend::Backend + $crate::types::HasSqlType<types::$Source>,
            $Target: $crate::types::FromSqlRow<types::$Source, DB>,
        {
            type Row = Self;

            fn build(row: Self::Row) -> Self {
                row
            }
        }
    )+}
}

#[doc(hidden)]
#[macro_export]
macro_rules! primitive_impls {
    ($Source:ident -> ($Target:ty, pg: ($oid:expr, $array_oid:expr), sqlite: ($tpe:ident))) => {
        #[cfg(feature = "sqlite")]
        impl types::HasSqlType<types::$Source> for $crate::sqlite::Sqlite {
            fn metadata() -> $crate::sqlite::SqliteType {
                $crate::sqlite::SqliteType::$tpe
            }
        }

        primitive_impls!($Source -> ($Target, pg: ($oid, $array_oid)));
    };

    ($Source:ident -> ($Target:ty, pg: ($oid:expr, $array_oid:expr))) => {
        primitive_impls!($Source -> (pg: ($oid, $array_oid)));
        primitive_impls!($Source -> $Target);
    };

    ($Source:ident -> (pg: ($oid:expr, $array_oid:expr))) => {
        #[cfg(feature = "postgres")]
        impl types::HasSqlType<types::$Source> for $crate::pg::Pg {
            fn metadata() -> $crate::pg::PgTypeMetadata {
                $crate::pg::PgTypeMetadata {
                    oid: $oid,
                    array_oid: $array_oid,
                }
            }
        }
    };

    ($Source:ident -> $Target:ty) => {
        primitive_impls!($Source);
        queryable_impls!($Source -> $Target,);
        expression_impls!($Source -> $Target,);
    };

    ($Source:ident) => {
        impl types::HasSqlType<types::$Source> for $crate::backend::Debug {
            fn metadata() {}
        }

        impl $crate::query_builder::QueryId for types::$Source {
            type QueryId = Self;

            fn has_static_query_id() -> bool {
                true
            }
        }

        impl types::NotNull for types::$Source {
        }
    }
}

macro_rules! debug_to_sql {
    ($sql_type:ty, $ty:ty) => {
        impl $crate::types::ToSql<$sql_type, $crate::backend::Debug> for $ty {
            fn to_sql<W: ::std::io::Write>(&self, _: &mut W) -> Result<$crate::types::IsNull, Box<::std::error::Error+Send+Sync>> {
                Ok($crate::types::IsNull::No)
            }
        }
    };
}

mod date_and_time;
pub mod floats;
mod integers;
pub mod option;
mod primitives;
mod tuples;
<|endoftext|>


<commit_before>extern crate num_bigint;
extern crate num_integer;
extern crate num_traits;

#[cfg(test)]
#[cfg(feature = "quickcheck")]
#[macro_use]
pub extern crate quickcheck;

#[cfg(test)]
#[cfg(feature = "quickcheck")]
mod quickchecks {
    use num_bigint::{BigInt, BigUint, Sign};
    use num_traits::{Num, One, Pow, Zero};
    use quickcheck::TestResult;

    quickcheck! {
        fn quickcheck_unsigned_eq_reflexive(a:BigUint) -> bool {
            a == a
        }
    }

    quickcheck! {
        fn quickcheck_signed_eq_reflexive(a:BigInt) -> bool {
            a == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_eq_symmetric(a:BigUint, b:BigUint) -> bool {
            if a == b {
                b == a
            } else {
                b != a
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_eq_symmetric(a:BigInt, b:BigInt) -> bool {
            if a == b {
                b == a
            } else {
                b != a
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_primitive(a: u128, b: u128) -> TestResult {
            let actual = BigUint::from(a) + BigUint::from(b);
            match a.checked_add(b) {
                None => TestResult::discard(),
                Some(expected) => TestResult::from_bool(BigUint::from(expected) == actual)
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_primitive(a: i128, b: i128) -> TestResult {
            let actual = BigInt::from(a) + BigInt::from(b);
            match a.checked_add(b) {
                None => TestResult::discard(),
                Some(expected) => TestResult::from_bool(BigInt::from(expected) == actual)
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_commutative(a: BigUint, b: BigUint) -> bool {
            a.clone() + b.clone() == b + a
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_commutative(a: BigInt, b: BigInt) -> bool {
            a.clone() + b.clone() == b + a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_zero(a: BigUint) -> bool {
            a.clone() == a + BigUint::zero()
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_zero(a: BigInt) -> bool {
            a.clone() == a + BigInt::zero()
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
            (a.clone() + b.clone()) + c.clone() == a + (b + c)
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
            (a.clone() + b.clone()) + c.clone() == a + (b + c)
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_primitive(a: u64, b: u64) -> bool {
            //maximum value of u64 means no overflow
            BigUint::from(a as u128 * b as u128) == BigUint::from(a) * BigUint::from(b)
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_primitive(a: i64, b: i64) -> bool {
            //maximum value of u64 means no overflow
            BigInt::from(a as i128 * b as i128) == BigInt::from(a) * BigInt::from(b)
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_zero(a: BigUint) -> bool {
            a * BigUint::zero() == BigUint::zero()
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_zero(a: BigInt) -> bool {
            a * BigInt::zero() == BigInt::zero()
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_one(a: BigUint) -> bool {
            a.clone() * BigUint::one() == a
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_one(a: BigInt) -> bool {
            a.clone() * BigInt::one() == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_commutative(a: BigUint, b: BigUint) -> bool {
            a.clone() * b.clone() == b * a
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_commutative(a: BigInt, b: BigInt) -> bool {
            a.clone() * b.clone() == b * a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
            (a.clone() * b.clone()) * c.clone() == a * (b * c)
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
            (a.clone() * b.clone()) * c.clone() == a * (b * c)
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_distributive(a: BigUint, b: BigUint, c: BigUint) -> bool {
            a.clone() * (b.clone() + c.clone()) == a.clone() * b + a * c
        }
    }

    quickcheck! {
        fn quickcheck_signed_distributive(a: BigInt, b: BigInt, c: BigInt) -> bool {
            a.clone() * (b.clone() + c.clone()) == a.clone() * b + a * c
        }
    }

    quickcheck! {
        ///Tests that exactly one of a<b a>b a=b is true
        fn quickcheck_unsigned_ge_le_eq_mut_exclusive(a: BigUint, b: BigUint) -> bool {
            let gt_lt_eq = vec![a > b, a < b, a == b];
            gt_lt_eq
                .iter()
                .fold(0, |acc, e| if *e { acc + 1 } else { acc })
                == 1
        }
    }

    quickcheck! {
        ///Tests that exactly one of a<b a>b a=b is true
        fn quickcheck_signed_ge_le_eq_mut_exclusive(a: BigInt, b: BigInt) -> bool {
            let gt_lt_eq = vec![a > b, a < b, a == b];
            gt_lt_eq
                .iter()
                .fold(0, |acc, e| if *e { acc + 1 } else { acc })
                == 1
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_sub_primitive(a: u128, b: u128) -> bool {
            if b < a {
                BigUint::from(a-b) == BigUint::from(a) - BigUint::from(b)
            } else {
                BigUint::from(b-a) == BigUint::from(b) - BigUint::from(a)
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_sub_primitive(a: i128, b: i128) -> bool {
            if b < a {
                BigInt::from(a-b) == BigInt::from(a) - BigInt::from(b)
            } else {
                BigInt::from(b-a) == BigInt::from(b) - BigInt::from(a)
            }
        }
    }

    quickcheck! {
        /// Tests correctness of subtraction assuming addition is correct
        fn quickcheck_unsigned_sub(a: BigUint, b: BigUint) -> bool {
            if b < a {
                a.clone() - b.clone() + b == a
            } else {
                b.clone() - a.clone() + a == b
            }
        }
    }

    quickcheck! {
        /// Tests correctness of subtraction assuming addition is correct
        fn quickcheck_signed_sub(a: BigInt, b: BigInt) -> bool {
            if b < a {
                a.clone() - b.clone() + b == a
            } else {
                b.clone() - a.clone() + a == b
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_div_primitive(a: u128, b: u128) -> TestResult {
            if b == 0 {
                TestResult::discard()
            } else {
                TestResult::from_bool(BigUint::from(a/b) == BigUint::from(a) / BigUint::from(b))
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_div_primitive(a: i128, b: i128) -> TestResult {
            if b == 0 {
                TestResult::discard()
            } else {
                TestResult::from_bool(BigInt::from(a/b) == BigInt::from(a) / BigInt::from(b))
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_pow_zero(a: BigUint) -> bool {
            a.pow(0_u32) == BigUint::one()
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_pow_one(a: BigUint) -> bool {
            a.pow(1_u32) == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_sqrt(a: BigUint) -> bool {
            (a.clone() * a.clone()).sqrt() == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_cbrt(a: BigUint) -> bool {
            (a.clone() * a.clone() * a.clone()).cbrt() == a
        }
    }

    quickcheck! {
        fn quickcheck_signed_cbrt(a: BigInt) -> bool {
            (a.clone() * a.clone() * a.clone()).cbrt() == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_conversion(a:BigUint) -> bool {
            let mut success = true;
            for radix in 2..=36 {
                let string = a.to_str_radix(radix);
                success &= a == BigUint::from_str_radix(&string, radix).unwrap()
            }
            success
        }
    }

    quickcheck! {
        fn quickcheck_signed_conversion(a:BigInt) -> bool {
            let mut success = true;
            for radix in 2..=36 {
                let string = a.to_str_radix(radix);
                success &= a == BigInt::from_str_radix(&string, radix).unwrap()
            }
            success
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_modpow(a:BigUint, exp:u8, modulus:BigUint) -> TestResult {
            if modulus.clone().is_zero() {
                TestResult::discard()
            } else {
                let expected = a.clone().pow(exp) % modulus.clone();
                let result = a.modpow(&BigUint::from(exp), &BigUint::from(modulus));
                TestResult::from_bool(expected == result)
            }
        }
    }
}
<commit_msg>remove extraneous import of Sign<commit_after>extern crate num_bigint;
extern crate num_integer;
extern crate num_traits;

#[cfg(test)]
#[cfg(feature = "quickcheck")]
#[macro_use]
pub extern crate quickcheck;

#[cfg(test)]
#[cfg(feature = "quickcheck")]
mod quickchecks {
    use num_bigint::{BigInt, BigUint};
    use num_traits::{Num, One, Pow, Zero};
    use quickcheck::TestResult;

    quickcheck! {
        fn quickcheck_unsigned_eq_reflexive(a:BigUint) -> bool {
            a == a
        }
    }

    quickcheck! {
        fn quickcheck_signed_eq_reflexive(a:BigInt) -> bool {
            a == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_eq_symmetric(a:BigUint, b:BigUint) -> bool {
            if a == b {
                b == a
            } else {
                b != a
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_eq_symmetric(a:BigInt, b:BigInt) -> bool {
            if a == b {
                b == a
            } else {
                b != a
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_primitive(a: u128, b: u128) -> TestResult {
            let actual = BigUint::from(a) + BigUint::from(b);
            match a.checked_add(b) {
                None => TestResult::discard(),
                Some(expected) => TestResult::from_bool(BigUint::from(expected) == actual)
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_primitive(a: i128, b: i128) -> TestResult {
            let actual = BigInt::from(a) + BigInt::from(b);
            match a.checked_add(b) {
                None => TestResult::discard(),
                Some(expected) => TestResult::from_bool(BigInt::from(expected) == actual)
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_commutative(a: BigUint, b: BigUint) -> bool {
            a.clone() + b.clone() == b + a
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_commutative(a: BigInt, b: BigInt) -> bool {
            a.clone() + b.clone() == b + a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_zero(a: BigUint) -> bool {
            a.clone() == a + BigUint::zero()
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_zero(a: BigInt) -> bool {
            a.clone() == a + BigInt::zero()
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_add_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
            (a.clone() + b.clone()) + c.clone() == a + (b + c)
        }
    }

    quickcheck! {
        fn quickcheck_signed_add_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
            (a.clone() + b.clone()) + c.clone() == a + (b + c)
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_primitive(a: u64, b: u64) -> bool {
            //maximum value of u64 means no overflow
            BigUint::from(a as u128 * b as u128) == BigUint::from(a) * BigUint::from(b)
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_primitive(a: i64, b: i64) -> bool {
            //maximum value of u64 means no overflow
            BigInt::from(a as i128 * b as i128) == BigInt::from(a) * BigInt::from(b)
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_zero(a: BigUint) -> bool {
            a * BigUint::zero() == BigUint::zero()
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_zero(a: BigInt) -> bool {
            a * BigInt::zero() == BigInt::zero()
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_one(a: BigUint) -> bool {
            a.clone() * BigUint::one() == a
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_one(a: BigInt) -> bool {
            a.clone() * BigInt::one() == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_commutative(a: BigUint, b: BigUint) -> bool {
            a.clone() * b.clone() == b * a
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_commutative(a: BigInt, b: BigInt) -> bool {
            a.clone() * b.clone() == b * a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_mul_associative(a: BigUint, b: BigUint, c: BigUint) -> bool {
            (a.clone() * b.clone()) * c.clone() == a * (b * c)
        }
    }

    quickcheck! {
        fn quickcheck_signed_mul_associative(a: BigInt, b: BigInt, c: BigInt) -> bool {
            (a.clone() * b.clone()) * c.clone() == a * (b * c)
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_distributive(a: BigUint, b: BigUint, c: BigUint) -> bool {
            a.clone() * (b.clone() + c.clone()) == a.clone() * b + a * c
        }
    }

    quickcheck! {
        fn quickcheck_signed_distributive(a: BigInt, b: BigInt, c: BigInt) -> bool {
            a.clone() * (b.clone() + c.clone()) == a.clone() * b + a * c
        }
    }

    quickcheck! {
        ///Tests that exactly one of a<b a>b a=b is true
        fn quickcheck_unsigned_ge_le_eq_mut_exclusive(a: BigUint, b: BigUint) -> bool {
            let gt_lt_eq = vec![a > b, a < b, a == b];
            gt_lt_eq
                .iter()
                .fold(0, |acc, e| if *e { acc + 1 } else { acc })
                == 1
        }
    }

    quickcheck! {
        ///Tests that exactly one of a<b a>b a=b is true
        fn quickcheck_signed_ge_le_eq_mut_exclusive(a: BigInt, b: BigInt) -> bool {
            let gt_lt_eq = vec![a > b, a < b, a == b];
            gt_lt_eq
                .iter()
                .fold(0, |acc, e| if *e { acc + 1 } else { acc })
                == 1
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_sub_primitive(a: u128, b: u128) -> bool {
            if b < a {
                BigUint::from(a-b) == BigUint::from(a) - BigUint::from(b)
            } else {
                BigUint::from(b-a) == BigUint::from(b) - BigUint::from(a)
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_sub_primitive(a: i128, b: i128) -> bool {
            if b < a {
                BigInt::from(a-b) == BigInt::from(a) - BigInt::from(b)
            } else {
                BigInt::from(b-a) == BigInt::from(b) - BigInt::from(a)
            }
        }
    }

    quickcheck! {
        /// Tests correctness of subtraction assuming addition is correct
        fn quickcheck_unsigned_sub(a: BigUint, b: BigUint) -> bool {
            if b < a {
                a.clone() - b.clone() + b == a
            } else {
                b.clone() - a.clone() + a == b
            }
        }
    }

    quickcheck! {
        /// Tests correctness of subtraction assuming addition is correct
        fn quickcheck_signed_sub(a: BigInt, b: BigInt) -> bool {
            if b < a {
                a.clone() - b.clone() + b == a
            } else {
                b.clone() - a.clone() + a == b
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_div_primitive(a: u128, b: u128) -> TestResult {
            if b == 0 {
                TestResult::discard()
            } else {
                TestResult::from_bool(BigUint::from(a/b) == BigUint::from(a) / BigUint::from(b))
            }
        }
    }

    quickcheck! {
        fn quickcheck_signed_div_primitive(a: i128, b: i128) -> TestResult {
            if b == 0 {
                TestResult::discard()
            } else {
                TestResult::from_bool(BigInt::from(a/b) == BigInt::from(a) / BigInt::from(b))
            }
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_pow_zero(a: BigUint) -> bool {
            a.pow(0_u32) == BigUint::one()
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_pow_one(a: BigUint) -> bool {
            a.pow(1_u32) == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_sqrt(a: BigUint) -> bool {
            (a.clone() * a.clone()).sqrt() == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_cbrt(a: BigUint) -> bool {
            (a.clone() * a.clone() * a.clone()).cbrt() == a
        }
    }

    quickcheck! {
        fn quickcheck_signed_cbrt(a: BigInt) -> bool {
            (a.clone() * a.clone() * a.clone()).cbrt() == a
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_conversion(a:BigUint) -> bool {
            let mut success = true;
            for radix in 2..=36 {
                let string = a.to_str_radix(radix);
                success &= a == BigUint::from_str_radix(&string, radix).unwrap()
            }
            success
        }
    }

    quickcheck! {
        fn quickcheck_signed_conversion(a:BigInt) -> bool {
            let mut success = true;
            for radix in 2..=36 {
                let string = a.to_str_radix(radix);
                success &= a == BigInt::from_str_radix(&string, radix).unwrap()
            }
            success
        }
    }

    quickcheck! {
        fn quickcheck_unsigned_modpow(a:BigUint, exp:u8, modulus:BigUint) -> TestResult {
            if modulus.clone().is_zero() {
                TestResult::discard()
            } else {
                let expected = a.clone().pow(exp) % modulus.clone();
                let result = a.modpow(&BigUint::from(exp), &BigUint::from(modulus));
                TestResult::from_bool(expected == result)
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::env;
use std::io::BufReader;
use std::io::BufRead;
use std::io::Write;
use std::collections::HashMap;
use std::collections::BTreeMap;
use std::collections::BTreeSet;

extern crate tools;
use tools::find_source_file;
use tools::file_format::analysis::{read_analysis, read_target, AnalysisKind};
use tools::config;

extern crate rustc_serialize;
use rustc_serialize::json::{Json, ToJson};

#[derive(Debug)]
struct SearchResult {
    lineno: u32,
    bounds: (u32, u32),
    line: String,
    context: String,
    contextsym: String,
}

impl ToJson for SearchResult {
    fn to_json(&self) -> Json {
        let (st, en) = self.bounds;
        let bounds = vec![st, en];

        let mut obj = BTreeMap::new();
        obj.insert("lno".to_string(), self.lineno.to_json());
        obj.insert("bounds".to_string(), bounds.to_json());
        obj.insert("line".to_string(), self.line.to_json());
        obj.insert("context".to_string(), self.context.to_json());
        obj.insert("contextsym".to_string(), self.contextsym.to_json());
        Json::Object(obj)
    }
}

fn split_scopes(id: &str) -> Vec<String> {
    let mut result = Vec::new();
    let mut start = 0;
    let mut argument_nesting = 0;
    for (index, m) in id.match_indices(|c| c == ':' || c == '<' || c == '>') {
        if m == ":" && argument_nesting == 0 {
            if start != index {
                result.push(id[start .. index].to_owned());
                start = index + 1;
            } else {
                start = index + 1;
            }
        } else if m == "<" {
            argument_nesting += 1;
        } else if m == ">" {
            argument_nesting -= 1;
        }
    }
    result.push(id[start ..].to_owned());
    return result;
}

fn main() {
    let args: Vec<_> = env::args().collect();

    let cfg = config::load(&args[1], false);

    let tree_name = &args[2];
    let tree_config = cfg.trees.get(tree_name).unwrap();

    let filenames_file = &args[3];

    let output_file = format!("{}/crossref", tree_config.paths.index_path);
    let jump_file = format!("{}/jumps", tree_config.paths.index_path);
    let id_file = format!("{}/identifiers", tree_config.paths.index_path);

    let mut table = HashMap::new();
    let mut pretty_table = HashMap::new();
    let mut id_table = HashMap::new();
    let mut jumps = Vec::new();

    {
        let mut process_file = |path: &str| {
            print!("File {}\n", path);

            let analysis_fname = format!("{}/analysis/{}", tree_config.paths.index_path, path);
            let analysis = read_analysis(&analysis_fname, &read_target);

            let source_fname = find_source_file(path, &tree_config.paths.files_path, &tree_config.paths.objdir_path);
            let source_file = match File::open(source_fname) {
                Ok(f) => f,
                Err(_) => {
                    println!("Unable to open source file");
                    return;
                },
            };
            let reader = BufReader::new(&source_file);
            let mut lines = Vec::new();
            for line in reader.lines() {
                match line {
                    Ok(l) => lines.push(l),
                    Err(_) => lines.push("".to_string()),
                }
            }

            for datum in analysis {
                for piece in datum.data {
                    let t1 = table.entry(piece.sym.to_owned()).or_insert(BTreeMap::new());
                    let t2 = t1.entry(piece.kind).or_insert(BTreeMap::new());
                    let t3 = t2.entry(path.to_owned()).or_insert(Vec::new());
                    let lineno = (datum.loc.lineno - 1) as usize;
                    if lineno >= lines.len() {
                        print!("Bad line number in file {} (line {})\n", path, lineno);
                        return;
                    }
                    let line = lines[lineno].clone();
                    let line_cut = line.trim_right();
                    let len = line_cut.len();
                    let line_cut = line_cut.trim_left();
                    let offset = (len - line_cut.len()) as u32;
                    let mut buf = String::new();
                    let mut i = 0;
                    for c in line_cut.chars() {
                        buf.push(c);
                        i += 1;
                        if i > 100 {
                            break;
                        }
                    }
                    t3.push(SearchResult {
                        lineno: datum.loc.lineno,
                        bounds: (datum.loc.col_start - offset, datum.loc.col_end - offset),
                        line: buf,
                        context: piece.context,
                        contextsym: piece.contextsym,
                    });

                    pretty_table.insert(piece.sym.to_owned(), piece.pretty.to_owned());

                    let ch = piece.sym.chars().nth(0).unwrap();
                    if !(ch >= '0' && ch <= '9') && !piece.sym.contains(' ') {
                        let t1 = id_table.entry(piece.pretty.to_owned()).or_insert(BTreeSet::new());
                        t1.insert(piece.sym.to_owned());
                    }
                }
            }
        };

        let f = File::open(filenames_file).unwrap();
        let file = BufReader::new(&f);
        for line in file.lines() {
            process_file(&line.unwrap());
        }
    }

    let mut outputf = File::create(output_file).unwrap();

    for (id, id_data) in table {
        let mut kindmap = BTreeMap::new();
        for (kind, kind_data) in &id_data {
            let mut result = Vec::new();
            for (path, results) in kind_data {
                let mut obj = BTreeMap::new();
                obj.insert("path".to_string(), path.to_json());
                obj.insert("lines".to_string(), results.to_json());
                result.push(Json::Object(obj));
            }
            let kindstr = match *kind {
                AnalysisKind::Use => "Uses",
                AnalysisKind::Def => "Definitions",
                AnalysisKind::Assign => "Assignments",
                AnalysisKind::Decl => "Declarations",
                AnalysisKind::Idl => "IDL",
            };
            kindmap.insert(kindstr.to_string(), Json::Array(result));
        }
        let kindmap = Json::Object(kindmap);

        let _ = outputf.write_all(format!("{}\n{}\n", id, kindmap.to_string()).as_bytes());

        if id_data.contains_key(&AnalysisKind::Def) {
            let defs = id_data.get(&AnalysisKind::Def).unwrap();
            if defs.len() == 1 {
                for (path, results) in defs {
                    if results.len() == 1 {
                        let mut v = Vec::new();
                        v.push(id.to_json());
                        v.push(path.to_json());
                        v.push(results[0].lineno.to_json());
                        let pretty = pretty_table.get(&id).unwrap();
                        v.push(pretty.to_json());
                        jumps.push(Json::Array(v));
                    }
                }
            }
        }
    }

    let mut jumpf = File::create(jump_file).unwrap();
    for jump in jumps {
        let _ = jumpf.write_all((jump.to_string() + "\n").as_bytes());
    }

    let mut idf = File::create(id_file).unwrap();
    for (id, syms) in id_table {
        for sym in syms {
            let components = split_scopes(&id);
            for i in 0..components.len() {
                let sub = &components[i..components.len()];
                let sub = sub.join("::");

                if !sub.is_empty() {
                    let line = format!("{} {}\n", sub, sym);
                    let _ = idf.write_all(line.as_bytes());
                }
            }
        }
    }
}
<commit_msg>Bug 1416876 - Make crossref output stable.<commit_after>use std::fs::File;
use std::env;
use std::io::BufReader;
use std::io::BufRead;
use std::io::Write;
use std::collections::HashMap;
use std::collections::BTreeMap;
use std::collections::BTreeSet;

extern crate tools;
use tools::find_source_file;
use tools::file_format::analysis::{read_analysis, read_target, AnalysisKind};
use tools::config;

extern crate rustc_serialize;
use rustc_serialize::json::{Json, ToJson};

#[derive(Debug)]
struct SearchResult {
    lineno: u32,
    bounds: (u32, u32),
    line: String,
    context: String,
    contextsym: String,
}

impl ToJson for SearchResult {
    fn to_json(&self) -> Json {
        let (st, en) = self.bounds;
        let bounds = vec![st, en];

        let mut obj = BTreeMap::new();
        obj.insert("lno".to_string(), self.lineno.to_json());
        obj.insert("bounds".to_string(), bounds.to_json());
        obj.insert("line".to_string(), self.line.to_json());
        obj.insert("context".to_string(), self.context.to_json());
        obj.insert("contextsym".to_string(), self.contextsym.to_json());
        Json::Object(obj)
    }
}

fn split_scopes(id: &str) -> Vec<String> {
    let mut result = Vec::new();
    let mut start = 0;
    let mut argument_nesting = 0;
    for (index, m) in id.match_indices(|c| c == ':' || c == '<' || c == '>') {
        if m == ":" && argument_nesting == 0 {
            if start != index {
                result.push(id[start .. index].to_owned());
                start = index + 1;
            } else {
                start = index + 1;
            }
        } else if m == "<" {
            argument_nesting += 1;
        } else if m == ">" {
            argument_nesting -= 1;
        }
    }
    result.push(id[start ..].to_owned());
    return result;
}

fn main() {
    let args: Vec<_> = env::args().collect();

    let cfg = config::load(&args[1], false);

    let tree_name = &args[2];
    let tree_config = cfg.trees.get(tree_name).unwrap();

    let filenames_file = &args[3];

    let output_file = format!("{}/crossref", tree_config.paths.index_path);
    let jump_file = format!("{}/jumps", tree_config.paths.index_path);
    let id_file = format!("{}/identifiers", tree_config.paths.index_path);

    let mut table = BTreeMap::new();
    let mut pretty_table = HashMap::new();
    let mut id_table = BTreeMap::new();
    let mut jumps = Vec::new();

    {
        let mut process_file = |path: &str| {
            print!("File {}\n", path);

            let analysis_fname = format!("{}/analysis/{}", tree_config.paths.index_path, path);
            let analysis = read_analysis(&analysis_fname, &read_target);

            let source_fname = find_source_file(path, &tree_config.paths.files_path, &tree_config.paths.objdir_path);
            let source_file = match File::open(source_fname) {
                Ok(f) => f,
                Err(_) => {
                    println!("Unable to open source file");
                    return;
                },
            };
            let reader = BufReader::new(&source_file);
            let mut lines = Vec::new();
            for line in reader.lines() {
                match line {
                    Ok(l) => lines.push(l),
                    Err(_) => lines.push("".to_string()),
                }
            }

            for datum in analysis {
                for piece in datum.data {
                    let t1 = table.entry(piece.sym.to_owned()).or_insert(BTreeMap::new());
                    let t2 = t1.entry(piece.kind).or_insert(BTreeMap::new());
                    let t3 = t2.entry(path.to_owned()).or_insert(Vec::new());
                    let lineno = (datum.loc.lineno - 1) as usize;
                    if lineno >= lines.len() {
                        print!("Bad line number in file {} (line {})\n", path, lineno);
                        return;
                    }
                    let line = lines[lineno].clone();
                    let line_cut = line.trim_right();
                    let len = line_cut.len();
                    let line_cut = line_cut.trim_left();
                    let offset = (len - line_cut.len()) as u32;
                    let mut buf = String::new();
                    let mut i = 0;
                    for c in line_cut.chars() {
                        buf.push(c);
                        i += 1;
                        if i > 100 {
                            break;
                        }
                    }
                    t3.push(SearchResult {
                        lineno: datum.loc.lineno,
                        bounds: (datum.loc.col_start - offset, datum.loc.col_end - offset),
                        line: buf,
                        context: piece.context,
                        contextsym: piece.contextsym,
                    });

                    pretty_table.insert(piece.sym.to_owned(), piece.pretty.to_owned());

                    let ch = piece.sym.chars().nth(0).unwrap();
                    if !(ch >= '0' && ch <= '9') && !piece.sym.contains(' ') {
                        let t1 = id_table.entry(piece.pretty.to_owned()).or_insert(BTreeSet::new());
                        t1.insert(piece.sym.to_owned());
                    }
                }
            }
        };

        let f = File::open(filenames_file).unwrap();
        let file = BufReader::new(&f);
        for line in file.lines() {
            process_file(&line.unwrap());
        }
    }

    let mut outputf = File::create(output_file).unwrap();

    for (id, id_data) in table {
        let mut kindmap = BTreeMap::new();
        for (kind, kind_data) in &id_data {
            let mut result = Vec::new();
            for (path, results) in kind_data {
                let mut obj = BTreeMap::new();
                obj.insert("path".to_string(), path.to_json());
                obj.insert("lines".to_string(), results.to_json());
                result.push(Json::Object(obj));
            }
            let kindstr = match *kind {
                AnalysisKind::Use => "Uses",
                AnalysisKind::Def => "Definitions",
                AnalysisKind::Assign => "Assignments",
                AnalysisKind::Decl => "Declarations",
                AnalysisKind::Idl => "IDL",
            };
            kindmap.insert(kindstr.to_string(), Json::Array(result));
        }
        let kindmap = Json::Object(kindmap);

        let _ = outputf.write_all(format!("{}\n{}\n", id, kindmap.to_string()).as_bytes());

        if id_data.contains_key(&AnalysisKind::Def) {
            let defs = id_data.get(&AnalysisKind::Def).unwrap();
            if defs.len() == 1 {
                for (path, results) in defs {
                    if results.len() == 1 {
                        let mut v = Vec::new();
                        v.push(id.to_json());
                        v.push(path.to_json());
                        v.push(results[0].lineno.to_json());
                        let pretty = pretty_table.get(&id).unwrap();
                        v.push(pretty.to_json());
                        jumps.push(Json::Array(v));
                    }
                }
            }
        }
    }

    let mut jumpf = File::create(jump_file).unwrap();
    for jump in jumps {
        let _ = jumpf.write_all((jump.to_string() + "\n").as_bytes());
    }

    let mut idf = File::create(id_file).unwrap();
    for (id, syms) in id_table {
        for sym in syms {
            let components = split_scopes(&id);
            for i in 0..components.len() {
                let sub = &components[i..components.len()];
                let sub = sub.join("::");

                if !sub.is_empty() {
                    let line = format!("{} {}\n", sub, sym);
                    let _ = idf.write_all(line.as_bytes());
                }
            }
        }
    }
}
<|endoftext|>