<commit_before>use std::io::prelude::*;

use term::terminfo::TermInfo;
use unicode_width::UnicodeWidthStr;

use std::sync::mpsc::{Receiver, Sender};
use std::borrow::Cow;
use std::collections::HashMap;
use std::ffi::OsString;

use std::sync::mpsc;
use std::io;
use std::env;
use std::thread;

use bis_c::{TermTrack, TermSize};
use error::StringError;
use search::SearchBase;

// TermControl contains utility funcitons for terminfo
struct TermControl {
    strings: HashMap<String, String>
}

#[derive(PartialEq)]
enum TermStack {
    Str(String),
    Int(isize),
    Bool(bool)
}

// our user interface instance
pub struct UI {
    // track is a guard, we never touch it manually
    #[allow(dead_code)]
    track: TermTrack,
    size: TermSize,
    query: Sender<String>,
    matches: Receiver<Vec<Cow<'static, str>>>,
    save_cursor: String,
    restore_cursor: String,
    clr_eos: String
}

impl Drop for UI {
    fn drop(&mut self) {
        debug!("Clearing screen on exit");
        let mut output = io::stdout();
        match write!(output, "{}", self.clr_eos) {
            Ok(_) => {
                trace!("Cleared screen successfully");
            },
            Err(e) => {
                error!("Failed to clear screen: {}", e);
            }
        }
        debug!("Flushing output");
        match output.flush() {
            Ok(_) => {
                trace!("Flushed screen successfully");
            },
            Err(e) => {
                error!("Failed to flush screen: {}", e);
            }
        }
    }
}

impl TermControl {
    pub fn create() -> Result<TermControl, StringError> {
        debug!("Getting terminal info");
        let info = match TermInfo::from_env() {
            Ok(info) => info,
            Err(e) => return Err(StringError::new("Failed to get TermInfo", Some(Box::new(e))))
        };

        trace!("Got terminfo: {:?}", info);

        let mut strings = HashMap::default();

        for (name, value) in info.strings.into_iter() {
            strings.insert(name, match String::from_utf8(value) {
                Ok(s) => s,
                Err(e) => return Err(StringError::new("Failed to convert value into an OsString", Some(Box::new(e))))
            });
        }

        // right now all we care about are the strings
        Ok(TermControl {
            strings: strings
        })
    }

    pub fn get_string<T: Borrow<String>>(&mut self, name: T, params: Vec<TermStack>) -> Option<String> {
        // only implement what we're actually using in the UI
        let sequence = match self.strings.get(name) {
            None => return None,
            Some(s) => {
                trace!("Matched string: {:?}", s);
                s.clone()
            }
        };

        let mut escaped = false;
        let mut stack: Vec<TermStack> = vec![];
        let mut result = String::default();
        let mut escape = String::default();

        for c in sequence.chars() {
            if !escaped {
                if c == '%' {
                    escaped = true;
                } else {
                    result.push(c);
                }
            } else if escape.is_empty() {
                if c == 'd' {
                    match stack.pop() {
                        Some(Int(c)) => {
                            result.push_all(format!("{}", c).as_ref());
                        },
                        _ => {}
                    }

                    escaped = false;
                } else if c == 'p' {
                    escaped.push('p');
                } else {
                    error!("Unknown escape character: {:?}", c);
                    escaped = false;
                }
            } else {
                if escape == "p" {

                }
            }
        }
    }
}

impl UI {
    pub fn create() -> Result<UI, StringError> {
        debug!("Getting terminfo");
        let mut info = match TermInfo::from_env() {
            Ok(t) => t,
            Err(e) => return Err(StringError::new("Failed to get terminfo",
                                                  Some(Box::new(e))))
        };

        // use terminfo to get control sequences
        debug!("Terminfo: {:?}", info);

        trace!("Getting save_cursor");
        let save_cursor = match info.strings.remove(&format!("sc")) {
            None => return Err(StringError::new("Terminfo did not contain save_cursor", None)),
            Some(item) => match String::from_utf8(item) {
                Err(e) => return Err(StringError::new("save_cursor was not valid utf-8", Some(Box::new(e)))),
                Ok(s) => {
                    trace!("save_cursor: {:?}", s.escape_default());
                    s
                }
            }
        };

        trace!("Getting restore_cursor");
        let restore_cursor = match info.strings.remove(&format!("rc")) {
            None => return Err(StringError::new("Terminfo did not contain restore_cursor", None)),
            Some(item) => match String::from_utf8(item) {
                Err(e) => return Err(StringError::new("restore_cursor was not valid utf-8", Some(Box::new(e)))),
                Ok(s) => {
                    trace!("restore_cursor: {:?}", s.escape_default());
                    s
                }
            }
        };

        trace!("Getting clr_eos");
        let clr_eos = match info.strings.remove(&format!("clr_eos")) {
            None => return Err(StringError::new("Terminfo did not contain clr_eos", None)),
            Some(item) => match String::from_utf8(item) {
                Err(e) => return Err(StringError::new("clr_eos was not valid utf-8", Some(Box::new(e)))),
                Ok(s) => {
                    trace!("restore_cursor: {:?}", s.escape_default());
                    s
                }
            }
        };

        let mut track = TermTrack::default();

        debug!("Getting terminal size");
        let size = match track.get_size() {
            Err(e) => return Err(StringError::new("Failed to get terminal size", Some(Box::new(e)))),
            Ok(s) => {
                trace!("Terminal size: {:?}", s);
                s
            }
        };

        debug!("Preparing terminal");
        match track.prepare() {
            Err(e) => return Err(StringError::new("Failed to prepare terminal", Some(Box::new(e)))),
            Ok(_) => {
                trace!("Terminal prepared successfully");
            }
        }

        debug!("Starting search thread");

        trace!("Creating thread primitives");
        let (query_tx, query_rx) = mpsc::channel();
        let (matches_tx, matches_rx) = mpsc::channel();

        trace!("Starting thread");
        thread::spawn(move|| {
            search_thread(query_rx, matches_tx);
        });

        debug!("Creating UI instance");
        let instance = UI {
            track: track,
            size: size,
            query: query_tx,
            matches: matches_rx,
            save_cursor: save_cursor,
            restore_cursor: restore_cursor,
            clr_eos: clr_eos
        };

        trace!("Instance creation successful");
        Ok(instance)
    }

    pub fn start(&mut self) -> Result<(), StringError> {
        // assume start on a new line
        // get handles for io
        let stdout = io::stdout();
        let stdin = io::stdin();
        let mut output = stdout.lock();
        let input = stdin.lock();

        let mut query = String::new();

        // draw our prompt, move ten lines down and then back up, and save the cursor
        debug!("Drawing prompt");
        match write!(output, "Match: {}", self.save_cursor) {
            Err(e) => return Err(StringError::new("Failed to draw prompt", Some(Box::new(e)))),
            Ok(_) => {
                trace!("Drew prompt successfully");
            }
        }

        // flush the output
        debug!("Flushing output");
        match output.flush() {
            Ok(_) => {
                trace!("Successfully flushed output");
            },
            Err(e) => {
                return Err(StringError::new("Failed to flush output", Some(Box::new(e))));
            }
        }

        // wait for input
        for maybe_chr in input.chars() {
            let chr = match maybe_chr {
                Err(e) => return Err(StringError::new("Failed to read character", Some(Box::new(e)))),
                Ok(c) => c
            };
            trace!("Got character: {:?}", chr);

            // push the character onto the query string
            query.push(chr);

            // draw the character, save the cursor position, clear the screen after us
            match write!(output, "{}{}{}", chr, self.save_cursor, self.clr_eos) {
                Err(e) => return Err(StringError::new("Failed to output character", Some(Box::new(e)))),
                Ok(_) => {
                    trace!("Outputted character successfully");
                }
            }

            // send the search thread our query
            debug!("Sending {} to search thread", &query);
            match self.query.send(query.clone()) {
                Ok(_) => {
                    trace!("Send successful");
                },
                Err(e) => {
                    return Err(StringError::new("Failed to send to search thread", Some(Box::new(e))));
                }
            }

            // wait for matches
            let matches = match self.matches.recv() {
                Ok(m) => m,
                Err(e) => {
                    return Err(StringError::new("Failed to read matches", Some(Box::new(e))));
                }
            };

            debug!("Got matches: {:?}", &matches);

            // draw the matches
            for item in matches.into_iter() {
                if UnicodeWidthStr::width(item.as_ref()) > self.size.cols {
                    let mut owned = item.into_owned();
                    while UnicodeWidthStr::width((&owned as &AsRef<str>).as_ref()) > self.size.cols {
                        // truncate long lines
                        owned.pop();
                    }
                    // draw the truncated item
                    match write!(output, "\n{}", owned) {
                        Err(e) => return Err(StringError::new("Failed to draw match", Some(Box::new(e)))),
                        Ok(_) => {
                            trace!("Drew match successfully");
                        }
                    }
                } else {
                    // draw the match after a newline
                    match write!(output, "\n{}", item) {
                        Err(e) => return Err(StringError::new("Failed to draw match", Some(Box::new(e)))),
                        Ok(_) => {
                            trace!("Drew match successfully");
                        }
                    }
                }
            }

            // restore the cursor
            match write!(output, "{}", self.restore_cursor) {
                Err(e) => return Err(StringError::new("Failed to restore cursor", Some(Box::new(e)))),
                Ok(_) => {
                    trace!("Restored cursor successfully");
                }
            }

            // flush the output
            match output.flush() {
                Ok(_) => {
                    trace!("Successfully flushed output");
                },
                Err(e) => {
                    return Err(StringError::new("Failed to flush output", Some(Box::new(e))));
                }
            }
        }

        // Return success
        Ok(())
    }
}

pub fn search_thread(query: Receiver<String>,
                     matches: Sender<Vec<Cow<'static, str>>>) {
    debug!("Starting query thread");

    debug!("Getting history path");
    let history_path = match env::var("HISTFILE") {
        Ok(p) => {
            trace!("Got history path: {:?}", p);
            p
        },
        Err(e) => panic!("Failed to get bash history file: {}", e)
    };

    let mut base = SearchBase::default();

    // read the history
    info!("Reading history");
    match base.read_history(history_path) {
        Ok(_) => {
            // success
        },
        Err(e) => {
            panic!("Failed to read history: {}", e)
        }
    }

    debug!("Starting query loop");

    loop {
        trace!("Waiting for a query");
        match query.recv() {
            Err(e) => {
                debug!("Search thread exiting: {}", e);
                break;
            },
            Ok(q) => {
                debug!("Got query: {:?}", q);
                let result = base.query(q);
                debug!("Got result: {:?}", result);
                match matches.send(result) {
                    Err(e) => {
                        debug!("Search thread exiting: {}", e);
                        break;
                    },
                    Ok(_) => {
                        trace!("Matches sent successfully");
                    }
                }
            }
        }
    }
}
<commit_msg>Improve UI<commit_after>use std::io::prelude::*;

use term::terminfo::TermInfo;
use unicode_width::UnicodeWidthStr;

use std::sync::mpsc::{Receiver, Sender};
use std::borrow::{Cow, Borrow};
use std::collections::HashMap;
use std::ffi::OsString;
use std::iter::FromIterator;

use std::sync::mpsc;
use std::io;
use std::env;
use std::thread;

use bis_c::{TermTrack, TermSize};
use error::StringError;
use search::SearchBase;

// TermControl contains utility funcitons for terminfo
#[derive(Debug)]
struct TermControl {
    strings: HashMap<String, String>
}

#[derive(PartialEq, Clone, Debug)]
enum TermStack {
    Str(String),
    Int(isize),
    Bool(bool)
}

// our user interface instance
pub struct UI {
    // track is a guard, we never touch it manually
    #[allow(dead_code)]
    track: TermTrack,
    size: TermSize,
    control: TermControl,
    query: Sender<String>,
    matches: Receiver<Vec<Cow<'static, str>>>,
}

impl Drop for UI {
    fn drop(&mut self) {
        debug!("Clearing screen on exit");
        let mut output = io::stdout();
        match write!(output, "{}", self.control.get_string("clr_eos".to_owned(), vec![]).unwrap_or(format!(""))) {
            Ok(_) => {
                trace!("Cleared screen successfully");
            },
            Err(e) => {
                error!("Failed to clear screen: {}", e);
            }
        }
        debug!("Flushing output");
        match output.flush() {
            Ok(_) => {
                trace!("Flushed screen successfully");
            },
            Err(e) => {
                error!("Failed to flush screen: {}", e);
            }
        }
    }
}

impl TermControl {
    pub fn create() -> Result<TermControl, StringError> {
        debug!("Getting terminal info");
        let info = match TermInfo::from_env() {
            Ok(info) => info,
            Err(e) => return Err(StringError::new("Failed to get TermInfo", Some(Box::new(e))))
        };

        trace!("Got terminfo: {:?}", info);

        let mut strings = HashMap::default();

        for (name, value) in info.strings.into_iter() {
            strings.insert(name, match String::from_utf8(value) {
                Ok(s) => s,
                Err(e) => return Err(StringError::new("Failed to convert value into an OsString", Some(Box::new(e))))
            });
        }

        // right now all we care about are the strings
        Ok(TermControl {
            strings: strings
        })
    }

    pub fn get_string<T: Borrow<String>>(&mut self, name: T, params: Vec<TermStack>) -> Option<String> {
        // only implement what we're actually using in the UI
        let sequence = match self.strings.get(name.borrow()) {
            None => {
                trace!("No match for string: {:?}", name.borrow());
                return None;
            },
            Some(s) => {
                trace!("Matched string: {:?}", s);
                s.clone()
            }
        };

        let mut escaped = false;
        let mut stack: Vec<TermStack> = vec![];
        let mut result = String::default();
        let mut escape = String::default();

        // only implement the sequences we care about

        for c in sequence.chars() {
            if !escaped {
                if c == '%' {
                    escaped = true;
                } else {
                    result.push(c);
                }
            } else if escape.is_empty() {
                if c == 'd' {
                    match stack.pop() {
                        Some(TermStack::Int(c)) => {
                            result.push_str(format!("{}", c).as_ref());
                        },
                        Some(o) => {
                            error!("Numeric print on non-numeric type: {:?}", o);
                        },
                        None => {
                            error!("Stack was empty on print");
                        }
                    }

                    escaped = false;
                } else if c == 'p' {
                    escape.push('p');
                } else {
                    error!("Unknown escape character: {:?}", c);
                    escaped = false;
                }
            } else {
                if escape == "p" {
                    match c.to_digit(10) {
                        Some(idx) => {
                            if idx != 0 {
                                match params.get(idx as usize - 1) {
                                    Some(item) => {
                                        stack.push(item.clone())
                                    },
                                    None => {
                                        error!("There was no parameter {}", idx);
                                    }
                                }
                            } else {
                                error!("Tried to print 0th paramater");
                            }
                        },
                        None => {
                            error!("Paramater number was not a digit");
                        }
                    }

                    escape.clear();
                    escaped = false;
                } else {
                    error!("Unknown escape sequence: {:?}", escape);
                    escape.clear();
                    escaped = false;
                }
            }
        }

        trace!("Returning result: {:?}", result);

        // return result
        Some(result)
    }
}

impl UI {
    pub fn create() -> Result<UI, StringError> {
        debug!("Creating TermControl");
        let control = try!(TermControl::create());

        trace!("Got TermControl: {:?}", control);

        let mut track = TermTrack::default();

        debug!("Getting terminal size");
        let size = match track.get_size() {
            Err(e) => return Err(StringError::new("Failed to get terminal size", Some(Box::new(e)))),
            Ok(s) => {
                trace!("Terminal size: {:?}", s);
                s
            }
        };

        debug!("Preparing terminal");
        match track.prepare() {
            Err(e) => return Err(StringError::new("Failed to prepare terminal", Some(Box::new(e)))),
            Ok(_) => {
                trace!("Terminal prepared successfully");
            }
        }

        debug!("Starting search thread");

        trace!("Creating thread primitives");
        let (query_tx, query_rx) = mpsc::channel();
        let (matches_tx, matches_rx) = mpsc::channel();

        trace!("Starting thread");
        thread::spawn(move|| {
            search_thread(query_rx, matches_tx);
        });

        debug!("Creating UI instance");
        let instance = UI {
            track: track,
            size: size,
            control: control,
            query: query_tx,
            matches: matches_rx,
        };

        trace!("Instance creation successful");
        Ok(instance)
    }

    pub fn start(&mut self) -> Result<(), StringError> {
        // assume start on a new line
        // get handles for io
        let stdout = io::stdout();
        let stdin = io::stdin();
        let mut output = stdout.lock();
        let input = stdin.lock();

        let mut query = String::new();

        // print a bunch of newlines, and then move back up
        match write!(output, "{}{}", String::from_iter(vec!['\n'; 10].into_iter()),
                     self.control.get_string("cuu".to_owned(), vec![TermStack::Int(10)]).unwrap_or(format!(""))) {
            Err(e) => return Err(StringError::new("Failed to create space", Some(Box::new(e)))),
            Ok(_) => {
                trace!("Successfully created space on terminal");
            }
        }

        // draw our prompt and save the cursor
        debug!("Drawing prompt");
        match write!(output, "Match: {}",
                     self.control.get_string("sc".to_owned(), vec![]).unwrap_or(format!(""))) {
            Err(e) => return Err(StringError::new("Failed to draw prompt", Some(Box::new(e)))),
            Ok(_) => {
                trace!("Drew prompt successfully");
            }
        }

        // flush the output
        debug!("Flushing output");
        match output.flush() {
            Ok(_) => {
                trace!("Successfully flushed output");
            },
            Err(e) => {
                return Err(StringError::new("Failed to flush output", Some(Box::new(e))));
            }
        }

        // wait for input
        for maybe_chr in input.chars() {
            let chr = match maybe_chr {
                Err(e) => return Err(StringError::new("Failed to read character", Some(Box::new(e)))),
                Ok(c) => c
            };
            trace!("Got character: {:?}", chr);

            // push the character onto the query string
            query.push(chr);

            // draw the character, save the cursor position, clear the screen after us
            match write!(output, "{}{}{}", chr,
                         self.control.get_string("sc".to_owned(), vec![]).unwrap_or(format!("")),
                         self.control.get_string("clr_eos".to_owned(), vec![]).unwrap_or(format!(""))) {
                Err(e) => return Err(StringError::new("Failed to output character", Some(Box::new(e)))),
                Ok(_) => {
                    trace!("Outputted character successfully");
                }
            }

            // send the search thread our query
            debug!("Sending {} to search thread", &query);
            match self.query.send(query.clone()) {
                Ok(_) => {
                    trace!("Send successful");
                },
                Err(e) => {
                    return Err(StringError::new("Failed to send to search thread", Some(Box::new(e))));
                }
            }

            // wait for matches
            let matches = match self.matches.recv() {
                Ok(m) => m,
                Err(e) => {
                    return Err(StringError::new("Failed to read matches", Some(Box::new(e))));
                }
            };

            debug!("Got matches: {:?}", &matches);

            // draw the matches
            for item in matches.into_iter() {
                if UnicodeWidthStr::width(item.as_ref()) > self.size.cols {
                    let mut owned = item.into_owned();
                    while UnicodeWidthStr::width((&owned as &AsRef<str>).as_ref()) > self.size.cols {
                        // truncate long lines
                        owned.pop();
                    }
                    // draw the truncated item
                    match write!(output, "\n{}", owned) {
                        Err(e) => return Err(StringError::new("Failed to draw match", Some(Box::new(e)))),
                        Ok(_) => {
                            trace!("Drew match successfully");
                        }
                    }
                } else {
                    // draw the match after a newline
                    match write!(output, "\n{}", item) {
                        Err(e) => return Err(StringError::new("Failed to draw match", Some(Box::new(e)))),
                        Ok(_) => {
                            trace!("Drew match successfully");
                        }
                    }
                }
            }

            // restore the cursor
            match write!(output, "{}", self.control.get_string("rc".to_owned(), vec![]).unwrap_or(format!(""))) {
                Err(e) => return Err(StringError::new("Failed to restore cursor", Some(Box::new(e)))),
                Ok(_) => {
                    trace!("Restored cursor successfully");
                }
            }

            // flush the output
            match output.flush() {
                Ok(_) => {
                    trace!("Successfully flushed output");
                },
                Err(e) => {
                    return Err(StringError::new("Failed to flush output", Some(Box::new(e))));
                }
            }
        }

        // Return success
        Ok(())
    }
}

pub fn search_thread(query: Receiver<String>,
                     matches: Sender<Vec<Cow<'static, str>>>) {
    debug!("Starting query thread");

    debug!("Getting history path");
    let history_path = match env::var("HISTFILE") {
        Ok(p) => {
            trace!("Got history path: {:?}", p);
            p
        },
        Err(e) => panic!("Failed to get bash history file: {}", e)
    };

    let mut base = SearchBase::default();

    // read the history
    info!("Reading history");
    match base.read_history(history_path) {
        Ok(_) => {
            // success
        },
        Err(e) => {
            panic!("Failed to read history: {}", e)
        }
    }

    debug!("Starting query loop");

    loop {
        trace!("Waiting for a query");
        match query.recv() {
            Err(e) => {
                debug!("Search thread exiting: {}", e);
                break;
            },
            Ok(q) => {
                debug!("Got query: {:?}", q);
                let result = base.query(q);
                debug!("Got result: {:?}", result);
                match matches.send(result) {
                    Err(e) => {
                        debug!("Search thread exiting: {}", e);
                        break;
                    },
                    Ok(_) => {
                        trace!("Matches sent successfully");
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use crate::task::JoinHandle;

cfg_rt_multi_thread! {
    /// Runs the provided blocking function on the current thread without
    /// blocking the executor.
    ///
    /// In general, issuing a blocking call or performing a lot of compute in a
    /// future without yielding is problematic, as it may prevent the executor
    /// from driving other tasks forward. Calling this function informs the
    /// executor that the currently executing task is about to block the thread,
    /// so the executor is able to hand off any other tasks it has to a new
    /// worker thread before that happens. See the [CPU-bound tasks and blocking
    /// code][blocking] section for more information.
    ///
    /// Be aware that although this function avoids starving other independently
    /// spawned tasks, any other code running concurrently in the same task will
    /// be suspended during the call to `block_in_place`. This can happen e.g.
    /// when using the [`join!`] macro. To avoid this issue, use
    /// [`spawn_blocking`] instead of `block_in_place`.
    ///
    /// Note that this function cannot be used within a [`current_thread`] runtime
    /// because in this case there are no other worker threads to hand off tasks
    /// to. On the other hand, calling the function outside a runtime is
    /// allowed. In this case, `block_in_place` just calls the provided closure
    /// normally.
    ///
    /// Code running behind `block_in_place` cannot be cancelled. When you shut
    /// down the executor, it will wait indefinitely for all blocking operations
    /// to finish. You can use [`shutdown_timeout`] to stop waiting for them
    /// after a certain timeout. Be aware that this will still not cancel the
    /// tasks — they are simply allowed to keep running after the method
    /// returns.
    ///
    /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code
    /// [`spawn_blocking`]: fn@crate::task::spawn_blocking
    /// [`join!`]: macro@join
    /// [`thread::spawn`]: fn@std::thread::spawn
    /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn docs() {
    /// task::block_in_place(move || {
    ///     // do some compute-heavy work or call synchronous code
    /// });
    /// # }
    /// ```
    ///
    /// Code running inside `block_in_place` may use `block_on` to reenter the
    /// async context.
    ///
    /// ```
    /// use tokio::task;
    /// use tokio::runtime::Handle;
    ///
    /// # async fn docs() {
    /// task::block_in_place(move || {
    ///     Handle::current().block_on(async move {
    ///         // do something async
    ///     });
    /// });
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// This function panics if called from a [`current_thread`] runtime.
    ///
    /// [`current_thread`]: fn@crate::runtime::Builder::new_current_thread
    pub fn block_in_place<F, R>(f: F) -> R
    where
        F: FnOnce() -> R,
    {
        crate::runtime::thread_pool::block_in_place(f)
    }
}

cfg_rt! {
    /// Runs the provided closure on a thread where blocking is acceptable.
    ///
    /// In general, issuing a blocking call or performing a lot of compute in a
    /// future without yielding is problematic, as it may prevent the executor from
    /// driving other futures forward. This function runs the provided closure on a
    /// thread dedicated to blocking operations. See the [CPU-bound tasks and
    /// blocking code][blocking] section for more information.
    ///
    /// Tokio will spawn more blocking threads when they are requested through this
    /// function until the upper limit configured on the [`Builder`] is reached.
    /// After reaching the upper limit, the tasks are put in a queue.
    /// The thread limit is very large by default, because `spawn_blocking` is often
    /// used for various kinds of IO operations that cannot be performed
    /// asynchronously.  When you run CPU-bound code using `spawn_blocking`, you
    /// should keep this large upper limit in mind. When running many CPU-bound
    /// computations, a semaphore or some other synchronization primitive should be
    /// used to limit the number of computation executed in parallel. Specialized
    /// CPU-bound executors, such as [rayon], may also be a good fit.
    ///
    /// This function is intended for non-async operations that eventually finish on
    /// their own. If you want to spawn an ordinary thread, you should use
    /// [`thread::spawn`] instead.
    ///
    /// Closures spawned using `spawn_blocking` cannot be cancelled. When you shut
    /// down the executor, it will wait indefinitely for all blocking operations to
    /// finish. You can use [`shutdown_timeout`] to stop waiting for them after a
    /// certain timeout. Be aware that this will still not cancel the tasks — they
    /// are simply allowed to keep running after the method returns.
    ///
    /// Note that if you are using the single threaded runtime, this function will
    /// still spawn additional threads for blocking operations. The basic
    /// scheduler's single thread is only used for asynchronous code.
    ///
    /// # Related APIs and patterns for bridging asynchronous and blocking code
    ///
    /// In simple cases, it is sufficient to have the closure accept input
    /// parameters at creation time and return a single value (or struct/tuple, etc.).
    ///
    /// For more complex situations in which it is desirable to stream data to or from
    /// the synchronous context, the [`mpsc channel`] has `blocking_send` and
    /// `blocking_recv` methods for use in non-async code such as the thread created
    /// by `spawn_blocking`.
    ///
    /// Another option is [`SyncIoBridge`] for cases where the synchronous context
    /// is operating on byte streams.  For example, you might use an asynchronous
    /// HTTP client such as [hyper] to fetch data, but perform complex parsing
    /// of the payload body using a library written for synchronous I/O.
    ///
    /// Finally, see also [Bridging with sync code][bridgesync] for discussions
    /// around the opposite case of using Tokio as part of a larger synchronous
    /// codebase.
    ///
    /// [`Builder`]: struct@crate::runtime::Builder
    /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code
    /// [rayon]: https://docs.rs/rayon
    /// [`mpsc channel`]: crate::sync::mpsc
    /// [`SyncIoBridge`]: https://docs.rs/tokio-util/0.6/tokio_util/io/struct.SyncIoBridge.html
    /// [hyper]: https://docs.rs/hyper
    /// [`thread::spawn`]: fn@std::thread::spawn
    /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout
    /// [bridgesync]: https://tokio.rs/tokio/topics/bridging
    ///
    /// # Examples
    ///
    /// Pass an input value and receive result of computation:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn docs() -> Result<(), Box<dyn std::error::Error>>{
    /// // Initial input
    /// let mut v = "Hello, ".to_string();
    /// let res = task::spawn_blocking(move || {
    ///     // Stand-in for compute-heavy work or using synchronous APIs
    ///     v.push_str("world");
    ///     // Pass ownership of the value back to the asynchronous context
    ///     v
    /// }).await?;
    ///
    /// // `res` is the value returned from the thread
    /// assert_eq!(res.as_str(), "Hello, world");
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// Use a channel:
    ///
    /// ```
    /// use tokio::task;
    /// use tokio::sync::mpsc;
    ///
    /// # async fn docs() {
    /// let (tx, mut rx) = mpsc::channel(2);
    /// let start = 5;
    /// let worker = task::spawn_blocking(move || {
    ///     for x in 0..10 {
    ///         // Stand in for complex computation
    ///         tx.blocking_send(start + x).unwrap();
    ///     }
    /// });
    ///
    /// let mut acc = 0;
    /// while let Some(v) = rx.recv().await {
    ///     acc += v;
    /// }
    /// assert_eq!(acc, 95);
    /// worker.await.unwrap();
    /// # }
    /// ```
    #[track_caller]
    pub fn spawn_blocking<F, R>(f: F) -> JoinHandle<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        crate::runtime::spawn_blocking(f)
    }
}
<commit_msg>task: expand on cancellation of `spawn_blocking` (#4811)<commit_after>use crate::task::JoinHandle;

cfg_rt_multi_thread! {
    /// Runs the provided blocking function on the current thread without
    /// blocking the executor.
    ///
    /// In general, issuing a blocking call or performing a lot of compute in a
    /// future without yielding is problematic, as it may prevent the executor
    /// from driving other tasks forward. Calling this function informs the
    /// executor that the currently executing task is about to block the thread,
    /// so the executor is able to hand off any other tasks it has to a new
    /// worker thread before that happens. See the [CPU-bound tasks and blocking
    /// code][blocking] section for more information.
    ///
    /// Be aware that although this function avoids starving other independently
    /// spawned tasks, any other code running concurrently in the same task will
    /// be suspended during the call to `block_in_place`. This can happen e.g.
    /// when using the [`join!`] macro. To avoid this issue, use
    /// [`spawn_blocking`] instead of `block_in_place`.
    ///
    /// Note that this function cannot be used within a [`current_thread`] runtime
    /// because in this case there are no other worker threads to hand off tasks
    /// to. On the other hand, calling the function outside a runtime is
    /// allowed. In this case, `block_in_place` just calls the provided closure
    /// normally.
    ///
    /// Code running behind `block_in_place` cannot be cancelled. When you shut
    /// down the executor, it will wait indefinitely for all blocking operations
    /// to finish. You can use [`shutdown_timeout`] to stop waiting for them
    /// after a certain timeout. Be aware that this will still not cancel the
    /// tasks — they are simply allowed to keep running after the method
    /// returns.
    ///
    /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code
    /// [`spawn_blocking`]: fn@crate::task::spawn_blocking
    /// [`join!`]: macro@join
    /// [`thread::spawn`]: fn@std::thread::spawn
    /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn docs() {
    /// task::block_in_place(move || {
    ///     // do some compute-heavy work or call synchronous code
    /// });
    /// # }
    /// ```
    ///
    /// Code running inside `block_in_place` may use `block_on` to reenter the
    /// async context.
    ///
    /// ```
    /// use tokio::task;
    /// use tokio::runtime::Handle;
    ///
    /// # async fn docs() {
    /// task::block_in_place(move || {
    ///     Handle::current().block_on(async move {
    ///         // do something async
    ///     });
    /// });
    /// # }
    /// ```
    ///
    /// # Panics
    ///
    /// This function panics if called from a [`current_thread`] runtime.
    ///
    /// [`current_thread`]: fn@crate::runtime::Builder::new_current_thread
    pub fn block_in_place<F, R>(f: F) -> R
    where
        F: FnOnce() -> R,
    {
        crate::runtime::thread_pool::block_in_place(f)
    }
}

cfg_rt! {
    /// Runs the provided closure on a thread where blocking is acceptable.
    ///
    /// In general, issuing a blocking call or performing a lot of compute in a
    /// future without yielding is problematic, as it may prevent the executor from
    /// driving other futures forward. This function runs the provided closure on a
    /// thread dedicated to blocking operations. See the [CPU-bound tasks and
    /// blocking code][blocking] section for more information.
    ///
    /// Tokio will spawn more blocking threads when they are requested through this
    /// function until the upper limit configured on the [`Builder`] is reached.
    /// After reaching the upper limit, the tasks are put in a queue.
    /// The thread limit is very large by default, because `spawn_blocking` is often
    /// used for various kinds of IO operations that cannot be performed
    /// asynchronously.  When you run CPU-bound code using `spawn_blocking`, you
    /// should keep this large upper limit in mind. When running many CPU-bound
    /// computations, a semaphore or some other synchronization primitive should be
    /// used to limit the number of computation executed in parallel. Specialized
    /// CPU-bound executors, such as [rayon], may also be a good fit.
    ///
    /// This function is intended for non-async operations that eventually finish on
    /// their own. If you want to spawn an ordinary thread, you should use
    /// [`thread::spawn`] instead.
    ///
    /// Closures spawned using `spawn_blocking` cannot be cancelled abruptly; there
    /// is no standard low level API to cause a thread to stop running.  However,
    /// a useful pattern is to pass some form of "cancellation token" into
    /// the thread.  This could be an [`AtomicBool`] that the task checks periodically.
    /// Another approach is to have the thread primarily read or write from a channel,
    /// and to exit when the channel closes; assuming the other side of the channel is dropped
    /// when cancellation occurs, this will cause the blocking task thread to exit
    /// soon after as well.
    ///
    /// When you shut down the executor, it will wait indefinitely for all blocking operations to
    /// finish. You can use [`shutdown_timeout`] to stop waiting for them after a
    /// certain timeout. Be aware that this will still not cancel the tasks — they
    /// are simply allowed to keep running after the method returns.  It is possible
    /// for a blocking task to be cancelled if it has not yet started running, but this
    /// is not guaranteed.
    ///
    /// Note that if you are using the single threaded runtime, this function will
    /// still spawn additional threads for blocking operations. The basic
    /// scheduler's single thread is only used for asynchronous code.
    ///
    /// # Related APIs and patterns for bridging asynchronous and blocking code
    ///
    /// In simple cases, it is sufficient to have the closure accept input
    /// parameters at creation time and return a single value (or struct/tuple, etc.).
    ///
    /// For more complex situations in which it is desirable to stream data to or from
    /// the synchronous context, the [`mpsc channel`] has `blocking_send` and
    /// `blocking_recv` methods for use in non-async code such as the thread created
    /// by `spawn_blocking`.
    ///
    /// Another option is [`SyncIoBridge`] for cases where the synchronous context
    /// is operating on byte streams.  For example, you might use an asynchronous
    /// HTTP client such as [hyper] to fetch data, but perform complex parsing
    /// of the payload body using a library written for synchronous I/O.
    ///
    /// Finally, see also [Bridging with sync code][bridgesync] for discussions
    /// around the opposite case of using Tokio as part of a larger synchronous
    /// codebase.
    ///
    /// [`Builder`]: struct@crate::runtime::Builder
    /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code
    /// [rayon]: https://docs.rs/rayon
    /// [`mpsc channel`]: crate::sync::mpsc
    /// [`SyncIoBridge`]: https://docs.rs/tokio-util/0.6/tokio_util/io/struct.SyncIoBridge.html
    /// [hyper]: https://docs.rs/hyper
    /// [`thread::spawn`]: fn@std::thread::spawn
    /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout
    /// [bridgesync]: https://tokio.rs/tokio/topics/bridging
    /// [`AtomicBool`]: struct@std::sync::atomic::AtomicBool
    ///
    /// # Examples
    ///
    /// Pass an input value and receive result of computation:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn docs() -> Result<(), Box<dyn std::error::Error>>{
    /// // Initial input
    /// let mut v = "Hello, ".to_string();
    /// let res = task::spawn_blocking(move || {
    ///     // Stand-in for compute-heavy work or using synchronous APIs
    ///     v.push_str("world");
    ///     // Pass ownership of the value back to the asynchronous context
    ///     v
    /// }).await?;
    ///
    /// // `res` is the value returned from the thread
    /// assert_eq!(res.as_str(), "Hello, world");
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// Use a channel:
    ///
    /// ```
    /// use tokio::task;
    /// use tokio::sync::mpsc;
    ///
    /// # async fn docs() {
    /// let (tx, mut rx) = mpsc::channel(2);
    /// let start = 5;
    /// let worker = task::spawn_blocking(move || {
    ///     for x in 0..10 {
    ///         // Stand in for complex computation
    ///         tx.blocking_send(start + x).unwrap();
    ///     }
    /// });
    ///
    /// let mut acc = 0;
    /// while let Some(v) = rx.recv().await {
    ///     acc += v;
    /// }
    /// assert_eq!(acc, 95);
    /// worker.await.unwrap();
    /// # }
    /// ```
    #[track_caller]
    pub fn spawn_blocking<F, R>(f: F) -> JoinHandle<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        crate::runtime::spawn_blocking(f)
    }
}
<|endoftext|>


<commit_before>//! The parser.

use std::borrow::Cow;

use crate::node::element::tag::{Tag, Type};
use crate::node::Attributes;

mod error;
mod reader;

pub use self::error::Error;

#[doc(hidden)]
pub use self::reader::Reader;

/// A parser.
pub struct Parser<'l> {
    #[allow(dead_code)]
    content: Cow<'l, str>,
    reader: Reader<'l>,
}

/// An event.
pub enum Event<'l> {
    /// An error.
    Error(Error),
    /// A tag.
    Tag(&'l str, Type, Attributes),
    /// A text.
    Text(&'l str),
    /// A comment.
    Comment,
    /// A declaration.
    Declaration,
    /// An instruction.
    Instruction,
}

/// A result.
pub type Result<T> = ::std::result::Result<T, Error>;

impl<'l> Parser<'l> {
    /// Create a parser.
    #[inline]
    pub fn new<T>(content: T) -> Self
    where
        T: Into<Cow<'l, str>>,
    {
        let content = content.into();
        let reader = unsafe { ::std::mem::transmute(Reader::new(&*content)) };
        Parser {
            content,
            reader,
        }
    }
}

macro_rules! raise(
    ($parser:expr, $($argument:tt)*) => (
        return Some(Event::Error(Error::new($parser.reader.position(), format!($($argument)*))));
    );
);

impl<'l> Iterator for Parser<'l> {
    type Item = Event<'l>;

    fn next(&mut self) -> Option<Self::Item> {
        let content = self.reader.capture(|reader| reader.consume_until_char('<'));
        if let Some(content) = content {
            return Some(Event::Text(content));
        }
        if !self.reader.consume_char('<') {
            return None;
        }
        let content = self.reader.capture(|reader| reader.consume_until_char('>'));
        if content.is_none() {
            raise!(self, "found an empty tag");
        }
        if !self.reader.consume_char('>') {
            raise!(self, "missing a closing angle bracket");
        }
        let content = content.unwrap();
        Some(if content.starts_with("!--") {
            Event::Comment
        } else if content.starts_with('!') {
            Event::Declaration
        } else if content.starts_with('?') {
            Event::Instruction
        } else {
            match Tag::parse(&content) {
                Ok(Tag(name, kind, attributes)) => Event::Tag(name, kind, attributes),
                Err(error) => Event::Error(error),
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::parser::{Event, Parser};

    #[test]
    fn next_tag() {
        macro_rules! test(
            ($content:expr, $value:expr) => ({
                let mut parser = Parser::new($content);
                match parser.next().unwrap() {
                    Event::Tag(value, _, _) => assert_eq!(value, $value),
                    _ => unreachable!(),
                }
            })
        );

        test!("<foo>", "foo");
        test!("<foo/>", "foo");
        test!("  <foo/>", "foo");
    }

    #[test]
    fn next_text() {
        macro_rules! test(
            ($content:expr, $value:expr) => ({
                let mut parser = Parser::new($content);
                match parser.next().unwrap() {
                    Event::Text(value) => assert_eq!(value, $value),
                    _ => unreachable!(),
                }
            })
        );

        test!("foo <bar>", "foo");
        test!("  foo<bar>", "foo");
        test!("foo> <bar>", "foo>");
    }
}
<commit_msg>Reformat source codes<commit_after>//! The parser.

use std::borrow::Cow;

use crate::node::element::tag::{Tag, Type};
use crate::node::Attributes;

mod error;
mod reader;

pub use self::error::Error;

#[doc(hidden)]
pub use self::reader::Reader;

/// A parser.
pub struct Parser<'l> {
    #[allow(dead_code)]
    content: Cow<'l, str>,
    reader: Reader<'l>,
}

/// An event.
pub enum Event<'l> {
    /// An error.
    Error(Error),
    /// A tag.
    Tag(&'l str, Type, Attributes),
    /// A text.
    Text(&'l str),
    /// A comment.
    Comment,
    /// A declaration.
    Declaration,
    /// An instruction.
    Instruction,
}

/// A result.
pub type Result<T> = ::std::result::Result<T, Error>;

impl<'l> Parser<'l> {
    /// Create a parser.
    #[inline]
    pub fn new<T>(content: T) -> Self
    where
        T: Into<Cow<'l, str>>,
    {
        let content = content.into();
        let reader = unsafe { ::std::mem::transmute(Reader::new(&*content)) };
        Parser { content, reader }
    }
}

macro_rules! raise(
    ($parser:expr, $($argument:tt)*) => (
        return Some(Event::Error(Error::new($parser.reader.position(), format!($($argument)*))));
    );
);

impl<'l> Iterator for Parser<'l> {
    type Item = Event<'l>;

    fn next(&mut self) -> Option<Self::Item> {
        let content = self.reader.capture(|reader| reader.consume_until_char('<'));
        if let Some(content) = content {
            return Some(Event::Text(content));
        }
        if !self.reader.consume_char('<') {
            return None;
        }
        let content = self.reader.capture(|reader| reader.consume_until_char('>'));
        if content.is_none() {
            raise!(self, "found an empty tag");
        }
        if !self.reader.consume_char('>') {
            raise!(self, "missing a closing angle bracket");
        }
        let content = content.unwrap();
        Some(if content.starts_with("!--") {
            Event::Comment
        } else if content.starts_with('!') {
            Event::Declaration
        } else if content.starts_with('?') {
            Event::Instruction
        } else {
            match Tag::parse(&content) {
                Ok(Tag(name, kind, attributes)) => Event::Tag(name, kind, attributes),
                Err(error) => Event::Error(error),
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::parser::{Event, Parser};

    #[test]
    fn next_tag() {
        macro_rules! test(
            ($content:expr, $value:expr) => ({
                let mut parser = Parser::new($content);
                match parser.next().unwrap() {
                    Event::Tag(value, _, _) => assert_eq!(value, $value),
                    _ => unreachable!(),
                }
            })
        );

        test!("<foo>", "foo");
        test!("<foo/>", "foo");
        test!("  <foo/>", "foo");
    }

    #[test]
    fn next_text() {
        macro_rules! test(
            ($content:expr, $value:expr) => ({
                let mut parser = Parser::new($content);
                match parser.next().unwrap() {
                    Event::Text(value) => assert_eq!(value, $value),
                    _ => unreachable!(),
                }
            })
        );

        test!("foo <bar>", "foo");
        test!("  foo<bar>", "foo");
        test!("foo> <bar>", "foo>");
    }
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Kata OS security coordinator fake support

extern crate alloc;
use alloc::fmt;
use alloc::string::{String, ToString};
use core::mem::size_of;
use core::ptr;
use hashbrown::HashMap;
use kata_memory_interface::kata_frame_alloc;
use kata_memory_interface::kata_frame_alloc_in_cnode;
use kata_memory_interface::kata_object_free_in_cnode;
use kata_memory_interface::kata_object_free_toplevel;
use kata_memory_interface::ObjDescBundle;
use kata_os_common::copyregion::CopyRegion;
use kata_os_common::cspace_slot::CSpaceSlot;
use kata_os_common::sel4_sys;
use kata_security_interface::*;
use kata_storage_interface::KeyValueData;
use log::trace;

use sel4_sys::seL4_PageBits;
use sel4_sys::seL4_Page_GetAddress;
use sel4_sys::seL4_Word;

const PAGE_SIZE: usize = 1 << seL4_PageBits;

extern "C" {
    // Regions for deep_copy work.
    static mut DEEP_COPY_SRC: [seL4_Word; PAGE_SIZE / size_of::<seL4_Word>()];
    static mut DEEP_COPY_DEST: [seL4_Word; PAGE_SIZE / size_of::<seL4_Word>()];
}

struct BundleData {
    pkg_contents: ObjDescBundle,
    pkg_size: usize,
    manifest: String,
    keys: HashMap<String, KeyValueData>,
}
impl BundleData {
    fn new(pkg_contents: &ObjDescBundle) -> Self {
        BundleData {
            pkg_contents: pkg_contents.clone(),
            pkg_size: pkg_contents.size_bytes(),
            manifest: String::from(
                r##"
# Comments like this
[Manifest]
BundleId=com.google.cerebra.hw.HelloWorld

[Binaries]
App=HelloWorldBin
Model=NeuralNetworkName

[Storage]
Required=1
"##,
            ),
            keys: HashMap::with_capacity(2),
        }
    }
}
impl Drop for BundleData {
    fn drop(&mut self) { let _ = kata_object_free_in_cnode(&self.pkg_contents); }
}

pub struct FakeSecurityCoordinator {
    bundles: HashMap<String, BundleData>,
}
impl Default for FakeSecurityCoordinator {
    fn default() -> Self { Self::new() }
}
impl FakeSecurityCoordinator {
    pub fn new() -> Self {
        FakeSecurityCoordinator {
            bundles: HashMap::with_capacity(2),
        }
    }

    fn get_bundle(&self, bundle_id: &str) -> Result<&BundleData, SecurityRequestError> {
        self.bundles
            .get(bundle_id)
            .map_or_else(|| Err(SecurityRequestError::SreBundleNotFound), Ok)
    }
    fn get_bundle_mut(&mut self, bundle_id: &str) -> Result<&mut BundleData, SecurityRequestError> {
        self.bundles
            .get_mut(bundle_id)
            .map_or_else(|| Err(SecurityRequestError::SreBundleNotFound), Ok)
    }
    fn remove_bundle(&mut self, bundle_id: &str) -> Result<(), SecurityRequestError> {
        self.bundles
            .remove(bundle_id)
            .map_or_else(|| Err(SecurityRequestError::SreBundleNotFound), |_| Ok(()))
    }
}
pub type KataSecurityCoordinatorInterface = FakeSecurityCoordinator;

// Returns a deep copy (including seL4 objects) of |src|. The container
// CNode is in the toplevel (allocated from the slot allocator).
fn deep_copy(src: &ObjDescBundle) -> ObjDescBundle {
    let dest = kata_frame_alloc_in_cnode(src.size_bytes()).expect("deep_copy:alloc");
    // Src top-level slot & copy region
    let src_slot = CSpaceSlot::new();
    let mut src_region = CopyRegion::new(unsafe { ptr::addr_of_mut!(DEEP_COPY_SRC[0]) }, PAGE_SIZE);
    // Dest top-level slot & copy region
    let dest_slot = CSpaceSlot::new();
    let mut dest_region =
        CopyRegion::new(unsafe { ptr::addr_of_mut!(DEEP_COPY_DEST[0]) }, PAGE_SIZE);
    for (src_cptr, dest_cptr) in src.cptr_iter().zip(dest.cptr_iter()) {
        // Map src & dest frames and copy data.
        src_slot
            .copy_to(src.cnode, src_cptr, src.depth)
            .and_then(|_| src_region.map(src_slot.slot))
            .expect("src_map");
        dest_slot
            .copy_to(dest.cnode, dest_cptr, dest.depth)
            .and_then(|_| dest_region.map(dest_slot.slot))
            .expect("dest_map");

        unsafe {
            ptr::copy_nonoverlapping(
                src_region.as_ref().as_ptr(),
                dest_region.as_mut().as_mut_ptr(),
                PAGE_SIZE,
            );
        }

        // Unmap & clear top-level slot required for mapping.
        src_region
            .unmap()
            .and_then(|_| src_slot.delete())
            .expect("src_unmap");
        dest_region
            .unmap()
            .and_then(|_| dest_slot.delete())
            .expect("dest_unmap");
    }
    dest
}

impl SecurityCoordinatorInterface for FakeSecurityCoordinator {
    fn install(&mut self, pkg_contents: &ObjDescBundle) -> Result<String, SecurityRequestError> {
        // TODO(sleffler): get bundle_id from the manifest; for now use the
        //    cnode's CPtr since it is unique wrt all installed packages
        let bundle_id = fmt::format(format_args!("fake.{}", pkg_contents.cnode));
        if self.bundles.contains_key(&bundle_id) {
            return Err(SecurityRequestError::SreDeleteFirst);
        }
        assert!(self
            .bundles
            .insert(bundle_id.clone(), BundleData::new(pkg_contents))
            .is_none());
        Ok(bundle_id)
    }
    fn uninstall(&mut self, bundle_id: &str) -> Result<(), SecurityRequestError> {
        self.remove_bundle(bundle_id)
    }
    fn size_buffer(&self, bundle_id: &str) -> Result<usize, SecurityRequestError> {
        let bundle = self.get_bundle(bundle_id)?;
        Ok(bundle.pkg_size) // TODO(sleffler): do better
    }
    fn get_manifest(&self, bundle_id: &str) -> Result<String, SecurityRequestError> {
        let bundle = self.get_bundle(bundle_id)?;
        // return &?
        Ok(bundle.manifest.clone())
    }
    fn load_application(&self, bundle_id: &str) -> Result<ObjDescBundle, SecurityRequestError> {
        let bundle_data = self.get_bundle(bundle_id)?;
        // Clone everything (struct + associated seL4 objects) so the
        // return is as though it was newly instantiated from flash.
        // XXX just return the package for now
        Ok(deep_copy(&bundle_data.pkg_contents))
    }
    fn load_model(
        &self,
        bundle_id: &str,
        _model_id: &str,
    ) -> Result<ObjDescBundle, SecurityRequestError> {
        let bundle_data = self.get_bundle(bundle_id)?;
        // TODO(sleffler): check model id
        // Clone everything (struct + associated seL4 objects) so the
        // return is as though it was newly instantiated from flash.
        // XXX just return the package for now
        Ok(deep_copy(&bundle_data.pkg_contents))
    }
    fn read_key(&self, bundle_id: &str, key: &str) -> Result<&KeyValueData, SecurityRequestError> {
        let bundle = self.get_bundle(bundle_id)?;
        bundle
            .keys
            .get(key)
            .ok_or(SecurityRequestError::SreKeyNotFound)
    }
    fn write_key(
        &mut self,
        bundle_id: &str,
        key: &str,
        value: &KeyValueData,
    ) -> Result<(), SecurityRequestError> {
        let bundle = self.get_bundle_mut(bundle_id)?;
        let _ = bundle.keys.insert(key.to_string(), *value);
        Ok(())
    }
    fn delete_key(&mut self, bundle_id: &str, key: &str) -> Result<(), SecurityRequestError> {
        let bundle = self.get_bundle_mut(bundle_id)?;
        // TODO(sleffler): error if no entry?
        let _ = bundle.keys.remove(key);
        Ok(())
    }

    fn test_mailbox(&mut self) -> Result<(), SecurityRequestError> {
        trace!("test_mailbox_command()");

        const MESSAGE_SIZE_DWORDS: usize = 17; // Just a random message size for testing.

        extern "C" {
            fn mailbox_api_send(paddr: u32, size: u32);
            fn mailbox_api_receive(paddr: *mut u32, size: *mut u32);
        }

        // Allocate a 4k page to serve as our message buffer.
        let frame_bundle =
            kata_frame_alloc(PAGE_SIZE).map_err(|_| SecurityRequestError::SreTestFailed)?;
        trace!("test_mailbox: Frame {:?}", frame_bundle);

        unsafe {
            // Map the message buffer into our copyregion so we can access it.
            // NB: re-use one of the deep_copy copyregions.
            let mut msg_region = CopyRegion::new(ptr::addr_of_mut!(DEEP_COPY_SRC[0]), PAGE_SIZE);
            msg_region
                .map(frame_bundle.objs[0].cptr)
                .map_err(|_| SecurityRequestError::SreTestFailed)?;

            let message_ptr = msg_region.as_word_mut();

            // Write to the message buffer through the copyregion.
            let offset_a = 0;
            let offset_b = MESSAGE_SIZE_DWORDS - 1;
            message_ptr[offset_a] = 0xDEADBEEF;
            message_ptr[offset_b] = 0xF00DCAFE;
            trace!(
                "test_mailbox: old buf contents  0x{:X} 0x{:X}",
                message_ptr[offset_a],
                message_ptr[offset_b]
            );

            // Send the _physical_ address of the message buffer to the security
            // core.
            let paddr = seL4_Page_GetAddress(frame_bundle.objs[0].cptr);
            mailbox_api_send(paddr.paddr as u32, (MESSAGE_SIZE_DWORDS * size_of::<u32>()) as u32);

            // Wait for the response to arrive.
            let mut response_paddr: u32 = 0;
            let mut response_size: u32 = 0;
            mailbox_api_receive(&mut response_paddr as *mut u32, &mut response_size as *mut u32);

            // The security core should have replaced the first and last dwords
            // with 0x12345678 and 0x87654321.
            trace!("test_mailbox: expected contents 0x12345678 0x87654321");
            trace!(
                "test_mailbox: new buf contents  0x{:X} 0x{:X}",
                message_ptr[offset_a],
                message_ptr[offset_b]
            );

            let dword_a = message_ptr[offset_a];
            let dword_b = message_ptr[offset_b];

            msg_region
                .unmap()
                .map_err(|_| SecurityRequestError::SreTestFailed)?;

            // Done, free the message buffer.
            kata_object_free_toplevel(&frame_bundle)
                .map_err(|_| SecurityRequestError::SreTestFailed)?;

            if dword_a != 0x12345678 || dword_b != 0x87654321 {
                return Err(SecurityRequestError::SreTestFailed);
            }
        }

        trace!("test_mailbox_command() done");
        Ok(())
    }
}
<commit_msg>kata-security-coordinator::fakeimpl: eliminate deep_copy panics<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Kata OS security coordinator fake support

extern crate alloc;
use alloc::fmt;
use alloc::string::{String, ToString};
use core::mem::size_of;
use core::ptr;
use hashbrown::HashMap;
use kata_memory_interface::kata_frame_alloc;
use kata_memory_interface::kata_frame_alloc_in_cnode;
use kata_memory_interface::kata_object_free_in_cnode;
use kata_memory_interface::kata_object_free_toplevel;
use kata_memory_interface::ObjDescBundle;
use kata_os_common::copyregion::CopyRegion;
use kata_os_common::cspace_slot::CSpaceSlot;
use kata_os_common::sel4_sys;
use kata_security_interface::*;
use kata_storage_interface::KeyValueData;
use log::trace;

use sel4_sys::seL4_Error;
use sel4_sys::seL4_PageBits;
use sel4_sys::seL4_Page_GetAddress;
use sel4_sys::seL4_Word;

const PAGE_SIZE: usize = 1 << seL4_PageBits;

extern "C" {
    // Regions for deep_copy work.
    static mut DEEP_COPY_SRC: [seL4_Word; PAGE_SIZE / size_of::<seL4_Word>()];
    static mut DEEP_COPY_DEST: [seL4_Word; PAGE_SIZE / size_of::<seL4_Word>()];
}

struct BundleData {
    pkg_contents: ObjDescBundle,
    pkg_size: usize,
    manifest: String,
    keys: HashMap<String, KeyValueData>,
}
impl BundleData {
    fn new(pkg_contents: &ObjDescBundle) -> Self {
        BundleData {
            pkg_contents: pkg_contents.clone(),
            pkg_size: pkg_contents.size_bytes(),
            manifest: String::from(
                r##"
# Comments like this
[Manifest]
BundleId=com.google.cerebra.hw.HelloWorld

[Binaries]
App=HelloWorldBin
Model=NeuralNetworkName

[Storage]
Required=1
"##,
            ),
            keys: HashMap::with_capacity(2),
        }
    }
}
impl Drop for BundleData {
    fn drop(&mut self) { let _ = kata_object_free_in_cnode(&self.pkg_contents); }
}

pub struct FakeSecurityCoordinator {
    bundles: HashMap<String, BundleData>,
}
impl Default for FakeSecurityCoordinator {
    fn default() -> Self { Self::new() }
}
impl FakeSecurityCoordinator {
    pub fn new() -> Self {
        FakeSecurityCoordinator {
            bundles: HashMap::with_capacity(2),
        }
    }

    fn get_bundle(&self, bundle_id: &str) -> Result<&BundleData, SecurityRequestError> {
        self.bundles
            .get(bundle_id)
            .map_or_else(|| Err(SecurityRequestError::SreBundleNotFound), Ok)
    }
    fn get_bundle_mut(&mut self, bundle_id: &str) -> Result<&mut BundleData, SecurityRequestError> {
        self.bundles
            .get_mut(bundle_id)
            .map_or_else(|| Err(SecurityRequestError::SreBundleNotFound), Ok)
    }
    fn remove_bundle(&mut self, bundle_id: &str) -> Result<(), SecurityRequestError> {
        self.bundles
            .remove(bundle_id)
            .map_or_else(|| Err(SecurityRequestError::SreBundleNotFound), |_| Ok(()))
    }
}
pub type KataSecurityCoordinatorInterface = FakeSecurityCoordinator;

// Returns a deep copy (including seL4 objects) of |src|. The container
// CNode is in the toplevel (allocated from the slot allocator).
fn deep_copy(src: &ObjDescBundle) -> Result<ObjDescBundle, seL4_Error> {
    let dest = kata_frame_alloc_in_cnode(src.size_bytes())
        .map_err(|_| seL4_Error::seL4_NotEnoughMemory)?; // TODO(sleffler) From mapping
    assert_eq!(src.count(), dest.count());
    // Src top-level slot & copy region
    let src_slot = CSpaceSlot::new();
    let mut src_region = CopyRegion::new(unsafe { ptr::addr_of_mut!(DEEP_COPY_SRC[0]) }, PAGE_SIZE);
    // Dest top-level slot & copy region
    let dest_slot = CSpaceSlot::new();
    let mut dest_region =
        CopyRegion::new(unsafe { ptr::addr_of_mut!(DEEP_COPY_DEST[0]) }, PAGE_SIZE);
    for (src_cptr, dest_cptr) in src.cptr_iter().zip(dest.cptr_iter()) {
        // Map src & dest frames and copy data.
        src_slot
            .copy_to(src.cnode, src_cptr, src.depth)
            .and_then(|_| src_region.map(src_slot.slot))?;
        dest_slot
            .copy_to(dest.cnode, dest_cptr, dest.depth)
            .and_then(|_| dest_region.map(dest_slot.slot))?;

        unsafe {
            ptr::copy_nonoverlapping(
                src_region.as_ref().as_ptr(),
                dest_region.as_mut().as_mut_ptr(),
                PAGE_SIZE,
            );
        }

        // Unmap & clear top-level slot required for mapping.
        src_region.unmap().and_then(|_| src_slot.delete())?;
        dest_region.unmap().and_then(|_| dest_slot.delete())?;
    }
    Ok(dest)
}

impl SecurityCoordinatorInterface for FakeSecurityCoordinator {
    fn install(&mut self, pkg_contents: &ObjDescBundle) -> Result<String, SecurityRequestError> {
        // TODO(sleffler): get bundle_id from the manifest; for now use the
        //    cnode's CPtr since it is unique wrt all installed packages
        let bundle_id = fmt::format(format_args!("fake.{}", pkg_contents.cnode));
        if self.bundles.contains_key(&bundle_id) {
            return Err(SecurityRequestError::SreDeleteFirst);
        }
        assert!(self
            .bundles
            .insert(bundle_id.clone(), BundleData::new(pkg_contents))
            .is_none());
        Ok(bundle_id)
    }
    fn uninstall(&mut self, bundle_id: &str) -> Result<(), SecurityRequestError> {
        self.remove_bundle(bundle_id)
    }
    fn size_buffer(&self, bundle_id: &str) -> Result<usize, SecurityRequestError> {
        let bundle = self.get_bundle(bundle_id)?;
        Ok(bundle.pkg_size) // TODO(sleffler): do better
    }
    fn get_manifest(&self, bundle_id: &str) -> Result<String, SecurityRequestError> {
        let bundle = self.get_bundle(bundle_id)?;
        // return &?
        Ok(bundle.manifest.clone())
    }
    fn load_application(&self, bundle_id: &str) -> Result<ObjDescBundle, SecurityRequestError> {
        let bundle_data = self.get_bundle(bundle_id)?;
        // Clone everything (struct + associated seL4 objects) so the
        // return is as though it was newly instantiated from flash.
        // XXX just return the package for now
        deep_copy(&bundle_data.pkg_contents)
            .map_err(|_| SecurityRequestError::SreLoadApplicationFailed)
    }
    fn load_model(
        &self,
        bundle_id: &str,
        _model_id: &str,
    ) -> Result<ObjDescBundle, SecurityRequestError> {
        let bundle_data = self.get_bundle(bundle_id)?;
        // TODO(sleffler): check model id
        // Clone everything (struct + associated seL4 objects) so the
        // return is as though it was newly instantiated from flash.
        // XXX just return the package for now
        deep_copy(&bundle_data.pkg_contents).map_err(|_| SecurityRequestError::SreLoadModelFailed)
    }
    fn read_key(&self, bundle_id: &str, key: &str) -> Result<&KeyValueData, SecurityRequestError> {
        let bundle = self.get_bundle(bundle_id)?;
        bundle
            .keys
            .get(key)
            .ok_or(SecurityRequestError::SreKeyNotFound)
    }
    fn write_key(
        &mut self,
        bundle_id: &str,
        key: &str,
        value: &KeyValueData,
    ) -> Result<(), SecurityRequestError> {
        let bundle = self.get_bundle_mut(bundle_id)?;
        let _ = bundle.keys.insert(key.to_string(), *value);
        Ok(())
    }
    fn delete_key(&mut self, bundle_id: &str, key: &str) -> Result<(), SecurityRequestError> {
        let bundle = self.get_bundle_mut(bundle_id)?;
        // TODO(sleffler): error if no entry?
        let _ = bundle.keys.remove(key);
        Ok(())
    }

    fn test_mailbox(&mut self) -> Result<(), SecurityRequestError> {
        trace!("test_mailbox_command()");

        const MESSAGE_SIZE_DWORDS: usize = 17; // Just a random message size for testing.

        extern "C" {
            fn mailbox_api_send(paddr: u32, size: u32);
            fn mailbox_api_receive(paddr: *mut u32, size: *mut u32);
        }

        // Allocate a 4k page to serve as our message buffer.
        let frame_bundle =
            kata_frame_alloc(PAGE_SIZE).map_err(|_| SecurityRequestError::SreTestFailed)?;
        trace!("test_mailbox: Frame {:?}", frame_bundle);

        unsafe {
            // Map the message buffer into our copyregion so we can access it.
            // NB: re-use one of the deep_copy copyregions.
            let mut msg_region = CopyRegion::new(ptr::addr_of_mut!(DEEP_COPY_SRC[0]), PAGE_SIZE);
            msg_region
                .map(frame_bundle.objs[0].cptr)
                .map_err(|_| SecurityRequestError::SreTestFailed)?;

            let message_ptr = msg_region.as_word_mut();

            // Write to the message buffer through the copyregion.
            let offset_a = 0;
            let offset_b = MESSAGE_SIZE_DWORDS - 1;
            message_ptr[offset_a] = 0xDEADBEEF;
            message_ptr[offset_b] = 0xF00DCAFE;
            trace!(
                "test_mailbox: old buf contents  0x{:X} 0x{:X}",
                message_ptr[offset_a],
                message_ptr[offset_b]
            );

            // Send the _physical_ address of the message buffer to the security
            // core.
            let paddr = seL4_Page_GetAddress(frame_bundle.objs[0].cptr);
            mailbox_api_send(paddr.paddr as u32, (MESSAGE_SIZE_DWORDS * size_of::<u32>()) as u32);

            // Wait for the response to arrive.
            let mut response_paddr: u32 = 0;
            let mut response_size: u32 = 0;
            mailbox_api_receive(&mut response_paddr as *mut u32, &mut response_size as *mut u32);

            // The security core should have replaced the first and last dwords
            // with 0x12345678 and 0x87654321.
            trace!("test_mailbox: expected contents 0x12345678 0x87654321");
            trace!(
                "test_mailbox: new buf contents  0x{:X} 0x{:X}",
                message_ptr[offset_a],
                message_ptr[offset_b]
            );

            let dword_a = message_ptr[offset_a];
            let dword_b = message_ptr[offset_b];

            msg_region
                .unmap()
                .map_err(|_| SecurityRequestError::SreTestFailed)?;

            // Done, free the message buffer.
            kata_object_free_toplevel(&frame_bundle)
                .map_err(|_| SecurityRequestError::SreTestFailed)?;

            if dword_a != 0x12345678 || dword_b != 0x87654321 {
                return Err(SecurityRequestError::SreTestFailed);
            }
        }

        trace!("test_mailbox_command() done");
        Ok(())
    }
}
<|endoftext|>


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

#[macro_use]
extern crate criterion;
use criterion::Criterion;

use rand::distributions::{Alphanumeric, Distribution, Standard};
use rand::prelude::random;
use rand::Rng;

use std::sync::Arc;

extern crate arrow;

use arrow::array::*;
use arrow::compute::{cast, take};
use arrow::datatypes::*;
use arrow::util::test_util::seedable_rng;

// cast array from specified primitive array type to desired data type
fn create_numeric<T>(size: usize) -> ArrayRef
where
    T: ArrowNumericType,
    Standard: Distribution<T::Native>,
    PrimitiveArray<T>: std::convert::From<Vec<T::Native>>,
{
    Arc::new(PrimitiveArray::<T>::from(vec![random::<T::Native>(); size])) as ArrayRef
}

fn create_strings(size: usize) -> ArrayRef {
    let v = (0..size)
        .map(|_| {
            seedable_rng()
                .sample_iter(&Alphanumeric)
                .take(5)
                .collect::<String>()
        })
        .collect::<Vec<_>>();

    Arc::new(StringArray::from(
        v.iter().map(|x| &**x).collect::<Vec<&str>>(),
    ))
}

fn create_random_index(size: usize) -> UInt32Array {
    let mut rng = seedable_rng();
    let ints = Int32Array::from(vec![rng.gen_range(-24i32, size as i32); size]);
    // cast to u32, conveniently marking negative values as nulls
    UInt32Array::from(
        cast(&(Arc::new(ints) as ArrayRef), &DataType::UInt32)
            .unwrap()
            .data(),
    )
}

fn bench_take(values: &ArrayRef, indices: &UInt32Array) {
    criterion::black_box(take(&values, &indices, None).unwrap());
}

fn add_benchmark(c: &mut Criterion) {
    let values = create_numeric::<Int32Type>(512);
    let indices = create_random_index(512);
    c.bench_function("take i32 512", |b| b.iter(|| bench_take(&values, &indices)));
    let values = create_numeric::<Int32Type>(1024);
    let indices = create_random_index(1024);
    c.bench_function("take i32 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = Arc::new(BooleanArray::from(vec![random::<bool>(); 512])) as ArrayRef;
    let indices = create_random_index(512);
    c.bench_function("take bool 512", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = Arc::new(BooleanArray::from(vec![random::<bool>(); 1024])) as ArrayRef;
    let indices = create_random_index(1024);
    c.bench_function("take bool 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = create_strings(512);
    let indices = create_random_index(512);
    c.bench_function("take str 512", |b| b.iter(|| bench_take(&values, &indices)));

    let values = create_strings(1024);
    let indices = create_random_index(1024);
    c.bench_function("take str 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });
}

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

#[macro_use]
extern crate criterion;
use criterion::Criterion;

use rand::distributions::{Alphanumeric, Distribution, Standard};
use rand::Rng;

use std::sync::Arc;

extern crate arrow;

use arrow::array::*;
use arrow::compute::take;
use arrow::datatypes::*;
use arrow::util::test_util::seedable_rng;

// cast array from specified primitive array type to desired data type
fn create_primitive<T>(size: usize) -> ArrayRef
where
    T: ArrowPrimitiveType,
    Standard: Distribution<T::Native>,
    PrimitiveArray<T>: std::convert::From<Vec<T::Native>>,
{
    let array: PrimitiveArray<T> = seedable_rng()
        .sample_iter(&Standard)
        .take(size)
        .map(|v| Some(v))
        .collect();

    Arc::new(array) as ArrayRef
}

fn create_strings(size: usize) -> ArrayRef {
    let v = (0..size)
        .map(|_| {
            seedable_rng()
                .sample_iter(&Alphanumeric)
                .take(5)
                .collect::<String>()
        })
        .collect::<Vec<_>>();

    Arc::new(StringArray::from(
        v.iter().map(|x| &**x).collect::<Vec<&str>>(),
    ))
}

fn create_random_index(size: usize, null_density: f32) -> UInt32Array {
    let mut rng = seedable_rng();
    let mut builder = UInt32Builder::new(size);
    for _ in 0..size {
        if rng.gen::<f32>() < null_density {
            let value = rng.gen_range::<u32, _, _>(0u32, size as u32);
            builder.append_value(value).unwrap();
        } else {
            builder.append_null().unwrap()
        }
    }
    builder.finish()
}

fn bench_take(values: &ArrayRef, indices: &UInt32Array) {
    criterion::black_box(take(&values, &indices, None).unwrap());
}

fn add_benchmark(c: &mut Criterion) {
    let values = create_primitive::<Int32Type>(512);
    let indices = create_random_index(512, 0.0);
    c.bench_function("take i32 512", |b| b.iter(|| bench_take(&values, &indices)));
    let values = create_primitive::<Int32Type>(1024);
    let indices = create_random_index(1024, 0.0);
    c.bench_function("take i32 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let indices = create_random_index(512, 0.5);
    c.bench_function("take i32 nulls 512", |b| {
        b.iter(|| bench_take(&values, &indices))
    });
    let values = create_primitive::<Int32Type>(1024);
    let indices = create_random_index(1024, 0.5);
    c.bench_function("take i32 nulls 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = create_primitive::<BooleanType>(512);
    let indices = create_random_index(512, 0.0);
    c.bench_function("take bool 512", |b| {
        b.iter(|| bench_take(&values, &indices))
    });
    let values = create_primitive::<BooleanType>(1024);
    let indices = create_random_index(1024, 0.0);
    c.bench_function("take bool 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = create_primitive::<BooleanType>(512);
    let indices = create_random_index(512, 0.5);
    c.bench_function("take bool nulls 512", |b| {
        b.iter(|| bench_take(&values, &indices))
    });
    let values = create_primitive::<BooleanType>(1024);
    let indices = create_random_index(1024, 0.5);
    c.bench_function("take bool nulls 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = create_strings(512);
    let indices = create_random_index(512, 0.0);
    c.bench_function("take str 512", |b| b.iter(|| bench_take(&values, &indices)));

    let values = create_strings(1024);
    let indices = create_random_index(1024, 0.0);
    c.bench_function("take str 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = create_strings(512);
    let indices = create_random_index(512, 0.5);
    c.bench_function("take str nulls 512", |b| {
        b.iter(|| bench_take(&values, &indices))
    });

    let values = create_strings(1024);
    let indices = create_random_index(1024, 0.5);
    c.bench_function("take str nulls 1024", |b| {
        b.iter(|| bench_take(&values, &indices))
    });
}

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


<commit_before>/// `mbe` (short for Macro By Example) crate contains code for handling
/// `macro_rules` macros. It uses `TokenTree` (from `ra_tt` package) as the
/// interface, although it contains some code to bridge `SyntaxNode`s and
/// `TokenTree`s as well!

mod parser;
mod mbe_expander;
mod syntax_bridge;
mod tt_iter;
mod subtree_source;

pub use tt::{Delimiter, Punct};

use crate::{
    parser::{parse_pattern, Op},
    tt_iter::TtIter,
};

#[derive(Debug, PartialEq, Eq)]
pub enum ParseError {
    Expected(String),
}

#[derive(Debug, PartialEq, Eq)]
pub enum ExpandError {
    NoMatchingRule,
    UnexpectedToken,
    BindingError(String),
    ConversionError,
    InvalidRepeat,
}

pub use crate::syntax_bridge::{
    ast_to_token_tree, syntax_node_to_token_tree, token_tree_to_expr, token_tree_to_items,
    token_tree_to_macro_stmts, token_tree_to_pat, token_tree_to_ty,
};

/// This struct contains AST for a single `macro_rules` definition. What might
/// be very confusing is that AST has almost exactly the same shape as
/// `tt::TokenTree`, but there's a crucial difference: in macro rules, `$ident`
/// and `$()*` have special meaning (see `Var` and `Repeat` data structures)
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MacroRules {
    pub(crate) rules: Vec<Rule>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct Rule {
    pub(crate) lhs: tt::Subtree,
    pub(crate) rhs: tt::Subtree,
}

impl MacroRules {
    pub fn parse(tt: &tt::Subtree) -> Result<MacroRules, ParseError> {
        let mut src = TtIter::new(tt);
        let mut rules = Vec::new();
        while src.len() > 0 {
            let rule = Rule::parse(&mut src)?;
            rules.push(rule);
            if let Err(()) = src.expect_char(';') {
                if src.len() > 0 {
                    return Err(ParseError::Expected("expected `:`".to_string()));
                }
                break;
            }
        }
        Ok(MacroRules { rules })
    }
    pub fn expand(&self, tt: &tt::Subtree) -> Result<tt::Subtree, ExpandError> {
        mbe_expander::expand(self, tt)
    }
}

impl Rule {
    fn parse(src: &mut TtIter) -> Result<Rule, ParseError> {
        let mut lhs = src
            .expect_subtree()
            .map_err(|()| ParseError::Expected("expected subtree".to_string()))?
            .clone();
        validate(&lhs)?;
        lhs.delimiter = tt::Delimiter::None;
        src.expect_char('=').map_err(|()| ParseError::Expected("expected `=`".to_string()))?;
        src.expect_char('>').map_err(|()| ParseError::Expected("expected `>`".to_string()))?;
        let mut rhs = src
            .expect_subtree()
            .map_err(|()| ParseError::Expected("expected subtree".to_string()))?
            .clone();
        rhs.delimiter = tt::Delimiter::None;
        Ok(crate::Rule { lhs, rhs })
    }
}

fn validate(pattern: &tt::Subtree) -> Result<(), ParseError> {
    for op in parse_pattern(pattern) {
        let op = match op {
            Ok(it) => it,
            Err(e) => {
                let msg = match e {
                    ExpandError::InvalidRepeat => "invalid repeat".to_string(),
                    _ => "invalid macro definition".to_string(),
                };
                return Err(ParseError::Expected(msg));
            }
        };
        match op {
            Op::TokenTree(tt::TokenTree::Subtree(subtree)) | Op::Repeat { subtree, .. } => {
                validate(subtree)?
            }
            _ => (),
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests;
<commit_msg>minor cleanup<commit_after>/// `mbe` (short for Macro By Example) crate contains code for handling
/// `macro_rules` macros. It uses `TokenTree` (from `ra_tt` package) as the
/// interface, although it contains some code to bridge `SyntaxNode`s and
/// `TokenTree`s as well!

mod parser;
mod mbe_expander;
mod syntax_bridge;
mod tt_iter;
mod subtree_source;

pub use tt::{Delimiter, Punct};

use crate::{
    parser::{parse_pattern, Op},
    tt_iter::TtIter,
};

#[derive(Debug, PartialEq, Eq)]
pub enum ParseError {
    Expected(String),
}

#[derive(Debug, PartialEq, Eq)]
pub enum ExpandError {
    NoMatchingRule,
    UnexpectedToken,
    BindingError(String),
    ConversionError,
    InvalidRepeat,
}

pub use crate::syntax_bridge::{
    ast_to_token_tree, syntax_node_to_token_tree, token_tree_to_expr, token_tree_to_items,
    token_tree_to_macro_stmts, token_tree_to_pat, token_tree_to_ty,
};

/// This struct contains AST for a single `macro_rules` definition. What might
/// be very confusing is that AST has almost exactly the same shape as
/// `tt::TokenTree`, but there's a crucial difference: in macro rules, `$ident`
/// and `$()*` have special meaning (see `Var` and `Repeat` data structures)
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MacroRules {
    pub(crate) rules: Vec<Rule>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct Rule {
    pub(crate) lhs: tt::Subtree,
    pub(crate) rhs: tt::Subtree,
}

impl MacroRules {
    pub fn parse(tt: &tt::Subtree) -> Result<MacroRules, ParseError> {
        // Note: this parsing can be implemented using mbe machinery itself, by
        // matching against `$($lhs:tt => $rhs:tt);*` pattern, but implementing
        // manually seems easier.
        let mut src = TtIter::new(tt);
        let mut rules = Vec::new();
        while src.len() > 0 {
            let rule = Rule::parse(&mut src)?;
            rules.push(rule);
            if let Err(()) = src.expect_char(';') {
                if src.len() > 0 {
                    return Err(ParseError::Expected("expected `:`".to_string()));
                }
                break;
            }
        }

        for rule in rules.iter() {
            validate(&rule.lhs)?;
        }

        Ok(MacroRules { rules })
    }
    pub fn expand(&self, tt: &tt::Subtree) -> Result<tt::Subtree, ExpandError> {
        mbe_expander::expand(self, tt)
    }
}

impl Rule {
    fn parse(src: &mut TtIter) -> Result<Rule, ParseError> {
        let mut lhs = src
            .expect_subtree()
            .map_err(|()| ParseError::Expected("expected subtree".to_string()))?
            .clone();
        lhs.delimiter = tt::Delimiter::None;
        src.expect_char('=').map_err(|()| ParseError::Expected("expected `=`".to_string()))?;
        src.expect_char('>').map_err(|()| ParseError::Expected("expected `>`".to_string()))?;
        let mut rhs = src
            .expect_subtree()
            .map_err(|()| ParseError::Expected("expected subtree".to_string()))?
            .clone();
        rhs.delimiter = tt::Delimiter::None;
        Ok(crate::Rule { lhs, rhs })
    }
}

fn validate(pattern: &tt::Subtree) -> Result<(), ParseError> {
    for op in parse_pattern(pattern) {
        let op = match op {
            Ok(it) => it,
            Err(e) => {
                let msg = match e {
                    ExpandError::InvalidRepeat => "invalid repeat".to_string(),
                    _ => "invalid macro definition".to_string(),
                };
                return Err(ParseError::Expected(msg));
            }
        };
        match op {
            Op::TokenTree(tt::TokenTree::Subtree(subtree)) | Op::Repeat { subtree, .. } => {
                validate(subtree)?
            }
            _ => (),
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests;
<|endoftext|>