<commit_before>//! `mdbook`'s low level rendering interface.
//!
//! # Note
//!
//! You usually don't need to work with this module directly. If you want to
//! implement your own backend, then check out the [For Developers] section of
//! the user guide.
//!
//! The definition for [RenderContext] may be useful though.
//!
//! [For Developers]: https://rust-lang-nursery.github.io/mdBook/lib/index.html
//! [RenderContext]: struct.RenderContext.html

pub use self::html_handlebars::HtmlHandlebars;

mod html_handlebars;

use std::io::Read;
use std::path::PathBuf;
use std::process::{Command, Stdio};
use serde_json;
use shlex::Shlex;

use errors::*;
use config::Config;
use book::Book;

/// An arbitrary `mdbook` backend.
///
/// Although it's quite possible for you to import `mdbook` as a library and
/// provide your own renderer, there are two main renderer implementations that
/// 99% of users will ever use:
///
/// - [HtmlHandlebars] - the built-in HTML renderer
/// - [CmdRenderer] - a generic renderer which shells out to a program to do the
///   actual rendering
///
/// [HtmlHandlebars]: struct.HtmlHandlebars.html
/// [CmdRenderer]: struct.CmdRenderer.html
pub trait Renderer {
    /// The `Renderer`'s name.
    fn name(&self) -> &str;

    /// Invoke the `Renderer`, passing in all the necessary information for
    /// describing a book.
    fn render(&self, ctx: &RenderContext) -> Result<()>;
}

/// The context provided to all renderers.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RenderContext {
    /// Which version of `mdbook` did this come from (as written in `mdbook`'s
    /// `Cargo.toml`). Useful if you know the renderer is only compatible with
    /// certain versions of `mdbook`.
    pub version: String,
    /// The book's root directory.
    pub root: PathBuf,
    /// A loaded representation of the book itself.
    pub book: Book,
    /// The loaded configuration file.
    pub config: Config,
    /// Where the renderer *must* put any build artefacts generated. To allow
    /// renderers to cache intermediate results, this directory is not
    /// guaranteed to be empty or even exist.
    pub destination: PathBuf,
}

impl RenderContext {
    /// Create a new `RenderContext`.
    pub(crate) fn new<P, Q>(root: P, book: Book, config: Config, destination: Q) -> RenderContext
    where
        P: Into<PathBuf>,
        Q: Into<PathBuf>,
    {
        RenderContext {
            book: book,
            config: config,
            version: env!("CARGO_PKG_VERSION").to_string(),
            root: root.into(),
            destination: destination.into(),
        }
    }

    /// Get the source directory's (absolute) path on disk.
    pub fn source_dir(&self) -> PathBuf {
        self.root.join(&self.config.book.src)
    }

    /// Load a `RenderContext` from its JSON representation.
    pub fn from_json<R: Read>(reader: R) -> Result<RenderContext> {
        serde_json::from_reader(reader).chain_err(|| "Unable to deserialize the `RenderContext`")
    }
}

/// A generic renderer which will shell out to an arbitrary executable.
///
/// # Rendering Protocol
///
/// When the renderer's `render()` method is invoked, `CmdRenderer` will spawn
/// the `cmd` as a subprocess. The `RenderContext` is passed to the subprocess
/// as a JSON string (using `serde_json`).
///
/// > **Note:** The command used doesn't necessarily need to be a single
/// > executable (i.e. `/path/to/renderer`). The `cmd` string lets you pass
/// > in command line arguments, so there's no reason why it couldn't be
/// > `python /path/to/renderer --from mdbook --to epub`.
///
/// Anything the subprocess writes to `stdin` or `stdout` will be passed through
/// to the user. While this gives the renderer maximum flexibility to output
/// whatever it wants, to avoid spamming users it is recommended to avoid
/// unnecessary output.
///
/// To help choose the appropriate output level, the `RUST_LOG` environment
/// variable will be passed through to the subprocess, if set.
///
/// If the subprocess wishes to indicate that rendering failed, it should exit
/// with a non-zero return code.
#[derive(Debug, Clone, PartialEq)]
pub struct CmdRenderer {
    name: String,
    cmd: String,
}

impl CmdRenderer {
    /// Create a new `CmdRenderer` which will invoke the provided `cmd` string.
    pub fn new(name: String, cmd: String) -> CmdRenderer {
        CmdRenderer { name, cmd }
    }

    fn compose_command(&self) -> Result<Command> {
        let mut words = Shlex::new(&self.cmd);
        let executable = match words.next() {
            Some(e) => e,
            None => bail!("Command string was empty"),
        };

        let mut cmd = Command::new(executable);

        for arg in words {
            cmd.arg(arg);
        }

        Ok(cmd)
    }
}

impl Renderer for CmdRenderer {
    fn name(&self) -> &str {
        &self.name
    }

    fn render(&self, ctx: &RenderContext) -> Result<()> {
        info!("Invoking the \"{}\" renderer", self.cmd);

        let mut child = self.compose_command()?
            .stdin(Stdio::piped())
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .current_dir(&ctx.destination)
            .spawn()
            .chain_err(|| "Unable to start the renderer")?;

        {
            let mut stdin = child.stdin.take().expect("Child has stdin");
            if let Err(e) = serde_json::to_writer(&mut stdin, &ctx) {
                // Looks like the backend hung up before we could finish
                // sending it the render context. Log the error and keep going
                warn!("Error writing the RenderContext to the backend, {}", e);
            }

            // explicitly close the `stdin` file handle
            drop(stdin);
        }

        let status = child
            .wait()
            .chain_err(|| "Error waiting for the backend to complete")?;

        trace!("{} exited with output: {:?}", self.cmd, status);

        if !status.success() {
            error!("Renderer exited with non-zero return code.");
            bail!("The \"{}\" renderer failed", self.cmd);
        } else {
            Ok(())
        }
    }
}
<commit_msg>Made sure the CmdRenderer always creates the destination dir<commit_after>//! `mdbook`'s low level rendering interface.
//!
//! # Note
//!
//! You usually don't need to work with this module directly. If you want to
//! implement your own backend, then check out the [For Developers] section of
//! the user guide.
//!
//! The definition for [RenderContext] may be useful though.
//!
//! [For Developers]: https://rust-lang-nursery.github.io/mdBook/lib/index.html
//! [RenderContext]: struct.RenderContext.html

pub use self::html_handlebars::HtmlHandlebars;

mod html_handlebars;

use std::fs;
use std::io::Read;
use std::path::PathBuf;
use std::process::{Command, Stdio};
use serde_json;
use shlex::Shlex;

use errors::*;
use config::Config;
use book::Book;

/// An arbitrary `mdbook` backend.
///
/// Although it's quite possible for you to import `mdbook` as a library and
/// provide your own renderer, there are two main renderer implementations that
/// 99% of users will ever use:
///
/// - [HtmlHandlebars] - the built-in HTML renderer
/// - [CmdRenderer] - a generic renderer which shells out to a program to do the
///   actual rendering
///
/// [HtmlHandlebars]: struct.HtmlHandlebars.html
/// [CmdRenderer]: struct.CmdRenderer.html
pub trait Renderer {
    /// The `Renderer`'s name.
    fn name(&self) -> &str;

    /// Invoke the `Renderer`, passing in all the necessary information for
    /// describing a book.
    fn render(&self, ctx: &RenderContext) -> Result<()>;
}

/// The context provided to all renderers.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RenderContext {
    /// Which version of `mdbook` did this come from (as written in `mdbook`'s
    /// `Cargo.toml`). Useful if you know the renderer is only compatible with
    /// certain versions of `mdbook`.
    pub version: String,
    /// The book's root directory.
    pub root: PathBuf,
    /// A loaded representation of the book itself.
    pub book: Book,
    /// The loaded configuration file.
    pub config: Config,
    /// Where the renderer *must* put any build artefacts generated. To allow
    /// renderers to cache intermediate results, this directory is not
    /// guaranteed to be empty or even exist.
    pub destination: PathBuf,
}

impl RenderContext {
    /// Create a new `RenderContext`.
    pub(crate) fn new<P, Q>(root: P, book: Book, config: Config, destination: Q) -> RenderContext
    where
        P: Into<PathBuf>,
        Q: Into<PathBuf>,
    {
        RenderContext {
            book: book,
            config: config,
            version: env!("CARGO_PKG_VERSION").to_string(),
            root: root.into(),
            destination: destination.into(),
        }
    }

    /// Get the source directory's (absolute) path on disk.
    pub fn source_dir(&self) -> PathBuf {
        self.root.join(&self.config.book.src)
    }

    /// Load a `RenderContext` from its JSON representation.
    pub fn from_json<R: Read>(reader: R) -> Result<RenderContext> {
        serde_json::from_reader(reader).chain_err(|| "Unable to deserialize the `RenderContext`")
    }
}

/// A generic renderer which will shell out to an arbitrary executable.
///
/// # Rendering Protocol
///
/// When the renderer's `render()` method is invoked, `CmdRenderer` will spawn
/// the `cmd` as a subprocess. The `RenderContext` is passed to the subprocess
/// as a JSON string (using `serde_json`).
///
/// > **Note:** The command used doesn't necessarily need to be a single
/// > executable (i.e. `/path/to/renderer`). The `cmd` string lets you pass
/// > in command line arguments, so there's no reason why it couldn't be
/// > `python /path/to/renderer --from mdbook --to epub`.
///
/// Anything the subprocess writes to `stdin` or `stdout` will be passed through
/// to the user. While this gives the renderer maximum flexibility to output
/// whatever it wants, to avoid spamming users it is recommended to avoid
/// unnecessary output.
///
/// To help choose the appropriate output level, the `RUST_LOG` environment
/// variable will be passed through to the subprocess, if set.
///
/// If the subprocess wishes to indicate that rendering failed, it should exit
/// with a non-zero return code.
#[derive(Debug, Clone, PartialEq)]
pub struct CmdRenderer {
    name: String,
    cmd: String,
}

impl CmdRenderer {
    /// Create a new `CmdRenderer` which will invoke the provided `cmd` string.
    pub fn new(name: String, cmd: String) -> CmdRenderer {
        CmdRenderer { name, cmd }
    }

    fn compose_command(&self) -> Result<Command> {
        let mut words = Shlex::new(&self.cmd);
        let executable = match words.next() {
            Some(e) => e,
            None => bail!("Command string was empty"),
        };

        let mut cmd = Command::new(executable);

        for arg in words {
            cmd.arg(arg);
        }

        Ok(cmd)
    }
}

impl Renderer for CmdRenderer {
    fn name(&self) -> &str {
        &self.name
    }

    fn render(&self, ctx: &RenderContext) -> Result<()> {
        info!("Invoking the \"{}\" renderer", self.cmd);

        let _ = fs::create_dir_all(&ctx.destination);

        let mut child = self.compose_command()?
            .stdin(Stdio::piped())
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .current_dir(&ctx.destination)
            .spawn()
            .chain_err(|| "Unable to start the renderer")?;

        {
            let mut stdin = child.stdin.take().expect("Child has stdin");
            if let Err(e) = serde_json::to_writer(&mut stdin, &ctx) {
                // Looks like the backend hung up before we could finish
                // sending it the render context. Log the error and keep going
                warn!("Error writing the RenderContext to the backend, {}", e);
            }

            // explicitly close the `stdin` file handle
            drop(stdin);
        }

        let status = child
            .wait()
            .chain_err(|| "Error waiting for the backend to complete")?;

        trace!("{} exited with output: {:?}", self.cmd, status);

        if !status.success() {
            error!("Renderer exited with non-zero return code.");
            bail!("The \"{}\" renderer failed", self.cmd);
        } else {
            Ok(())
        }
    }
}
<|endoftext|>


<commit_before>use std;
use std::sync::Arc;
use std::collections::VecDeque;

use rope::Rope;
use str_utils::{is_grapheme_boundary, nearest_internal_grapheme_boundary, next_grapheme_boundary,
                prev_grapheme_boundary};
use tree::{Node, NodeChildren, NodeText, MAX_BYTES, MAX_CHILDREN};

/// An incremental `Rope` builder.
///
/// `RopeBuilder` is used to efficiently build `Rope`s from sequences
/// of text chunks.  It is useful for situations such as:
///
/// - Creating a rope from a large text file without pre-loading the
///   entire contents of the file into memory first (but see
///   `Rope::from_reader()` which uses `RopeBuilder` internally for
///   precisely this use-case).
/// - Creating a rope from a streaming data source.
/// - Loading a non-utf8 text source into a rope, doing the encoding
///   conversion incrementally as you go.
///
/// Unlike repeatedly calling `Rope::insert()` on the end of a rope,
/// this API runs in time linear to the amount of data fed to it.  It
/// also creates more memory-compact ropes.
///
/// (The converse of this API is the [`Chunks`](iter/struct.Chunks.html)
/// iterator.)
///
/// # Example
/// ```
/// # use ropey::RopeBuilder;
/// #
/// let mut builder = RopeBuilder::new();
///
/// builder.append("Hello ");
/// builder.append("world!\n");
/// builder.append("How's ");
/// builder.append("it goin");
/// builder.append("g?");
///
/// let rope = builder.finish();
///
/// assert_eq!(rope, "Hello world!\nHow's it going?");
/// ```
#[derive(Debug, Clone)]
pub struct RopeBuilder {
    stack: VecDeque<Node>,
    buffer: String,
}

impl RopeBuilder {
    /// Creates a new RopeBuilder, ready for input.
    pub fn new() -> RopeBuilder {
        RopeBuilder {
            stack: {
                let mut stack = VecDeque::with_capacity(8);
                stack.push_back(Node::new());
                stack
            },
            buffer: String::with_capacity(MAX_BYTES),
        }
    }

    /// Appends `chunk` to the end of the in-progress `Rope`.
    ///
    /// This method is called repeatedly to incrementally build up a
    /// `Rope`.  The passed text chunk can be as large or small as
    /// desired, but larger chunks are more efficient.
    ///
    /// `chunk` must be valid utf8 text.
    pub fn append(&mut self, chunk: &str) {
        let mut chunk = chunk;

        // Repeatedly chop text off the end of the input, creating
        // leaf nodes out of them and appending them to the tree.
        while !chunk.is_empty() {
            // Get the text for the next leaf
            let (leaf_text, remainder) = self.get_next_leaf_text(chunk);
            chunk = remainder;

            // Append the leaf to the rope
            match leaf_text {
                NextText::None => break,
                NextText::UseBuffer => {
                    let string = NodeText::from_str(&self.buffer);
                    self.append_leaf_node(Node::Leaf(string));
                    self.buffer.clear();
                }
                NextText::String(s) => {
                    self.append_leaf_node(Node::Leaf(NodeText::from_str(s)));
                }
            }
        }
    }

    /// Finishes the build, and returns the `Rope`.
    ///
    /// Note: this consumes the builder.  If you want to continue building
    /// other ropes with the same prefix, you can clone the builder before
    /// calling `finish()`.
    pub fn finish(mut self) -> Rope {
        // Append the last leaf
        if !self.buffer.is_empty() {
            let string = NodeText::from_str(&self.buffer);
            self.append_leaf_node(Node::Leaf(string));
        }

        // Zip up all the remaining nodes on the stack
        let mut stack_idx = self.stack.len() - 1;
        while stack_idx >= 1 {
            let node = self.stack.pop_back().unwrap();
            if let Node::Internal(ref mut children) = self.stack[stack_idx - 1] {
                children.push((node.text_info(), Arc::new(node)));
            } else {
                unreachable!();
            }
            stack_idx -= 1;
        }

        // Get root and fix any right-side nodes with too few children.
        let mut root = self.stack.pop_back().unwrap();
        root.zip_fix_right();

        // Crate the rope, make sure it's well-formed, and return it.
        let mut rope = Rope {
            root: Arc::new(root),
        };
        rope.pull_up_singular_nodes();
        return rope;
    }

    //-----------------------------------------------------------------

    // Returns (next_leaf_text, remaining_text)
    fn get_next_leaf_text<'a>(&mut self, text: &'a str) -> (NextText<'a>, &'a str) {
        if self.buffer.is_empty() {
            if text.len() > MAX_BYTES {
                // Simplest case: just chop off the end of `text`
                let split_idx = find_good_split_idx(text, MAX_BYTES);
                if split_idx == 0 || split_idx == text.len() {
                    self.buffer.push_str(text);
                    return (NextText::None, "");
                } else {
                    return (NextText::String(&text[..split_idx]), &text[split_idx..]);
                }
            } else {
                self.buffer.push_str(text);
                return (NextText::None, "");
            }
        } else if (text.len() + self.buffer.len()) > MAX_BYTES {
            let split_idx = if self.buffer.len() < MAX_BYTES {
                nearest_internal_grapheme_boundary(text, MAX_BYTES - self.buffer.len())
            } else {
                nearest_internal_grapheme_boundary(text, 0)
            };

            if split_idx == 0 || split_idx == text.len() {
                self.buffer.push_str(text);
                return (NextText::None, "");
            } else {
                self.buffer.push_str(&text[..split_idx]);
                return (NextText::UseBuffer, &text[split_idx..]);
            }
        } else {
            self.buffer.push_str(text);
            return (NextText::None, "");
        }
    }

    fn append_leaf_node(&mut self, leaf: Node) {
        let last = self.stack.pop_back().unwrap();
        match last {
            Node::Leaf(_) => {
                if last.leaf_text().is_empty() {
                    self.stack.push_back(leaf);
                } else {
                    let mut children = NodeChildren::new();
                    children.push((last.text_info(), Arc::new(last)));
                    children.push((leaf.text_info(), Arc::new(leaf)));
                    self.stack.push_back(Node::Internal(children));
                }
            }

            Node::Internal(_) => {
                self.stack.push_back(last);
                let mut left = leaf;
                let mut stack_idx = (self.stack.len() - 1) as isize;
                loop {
                    if stack_idx < 0 {
                        // We're above the root, so do a root split.
                        let mut children = NodeChildren::new();
                        children.push((left.text_info(), Arc::new(left)));
                        self.stack.push_front(Node::Internal(children));
                        break;
                    } else if self.stack[stack_idx as usize].child_count() < (MAX_CHILDREN - 1) {
                        // There's room to add a child, so do that.
                        self.stack[stack_idx as usize]
                            .children()
                            .push((left.text_info(), Arc::new(left)));
                        break;
                    } else {
                        // Not enough room to fit a child, so split.
                        left = Node::Internal(
                            self.stack[stack_idx as usize]
                                .children()
                                .push_split((left.text_info(), Arc::new(left))),
                        );
                        std::mem::swap(&mut left, &mut self.stack[stack_idx as usize]);
                        stack_idx -= 1;
                    }
                }
            }
        }
    }
}

enum NextText<'a> {
    None,
    UseBuffer,
    String(&'a str),
}

fn find_good_split_idx(text: &str, idx: usize) -> usize {
    if is_grapheme_boundary(text, idx) {
        idx
    } else {
        let prev = prev_grapheme_boundary(text, idx);
        let next = next_grapheme_boundary(text, idx);
        if prev > 0 {
            prev
        } else {
            next
        }
    }
}
<commit_msg>Changed RopeBuilder to use a SmallVec of Arc<Node> for its stack.<commit_after>use std;
use std::sync::Arc;
use smallvec::SmallVec;

use rope::Rope;
use str_utils::{is_grapheme_boundary, nearest_internal_grapheme_boundary, next_grapheme_boundary,
                prev_grapheme_boundary};
use tree::{Node, NodeChildren, NodeText, MAX_BYTES, MAX_CHILDREN};

/// An incremental `Rope` builder.
///
/// `RopeBuilder` is used to efficiently build `Rope`s from sequences
/// of text chunks.  It is useful for situations such as:
///
/// - Creating a rope from a large text file without pre-loading the
///   entire contents of the file into memory first (but see
///   `Rope::from_reader()` which uses `RopeBuilder` internally for
///   precisely this use-case).
/// - Creating a rope from a streaming data source.
/// - Loading a non-utf8 text source into a rope, doing the encoding
///   conversion incrementally as you go.
///
/// Unlike repeatedly calling `Rope::insert()` on the end of a rope,
/// this API runs in time linear to the amount of data fed to it.  It
/// also creates more memory-compact ropes.
///
/// (The converse of this API is the [`Chunks`](iter/struct.Chunks.html)
/// iterator.)
///
/// # Example
/// ```
/// # use ropey::RopeBuilder;
/// #
/// let mut builder = RopeBuilder::new();
///
/// builder.append("Hello ");
/// builder.append("world!\n");
/// builder.append("How's ");
/// builder.append("it goin");
/// builder.append("g?");
///
/// let rope = builder.finish();
///
/// assert_eq!(rope, "Hello world!\nHow's it going?");
/// ```
#[derive(Debug, Clone)]
pub struct RopeBuilder {
    stack: SmallVec<[Arc<Node>; 4]>,
    buffer: String,
}

impl RopeBuilder {
    /// Creates a new RopeBuilder, ready for input.
    pub fn new() -> RopeBuilder {
        RopeBuilder {
            stack: {
                let mut stack = SmallVec::new();
                stack.push(Arc::new(Node::new()));
                stack
            },
            buffer: String::new(),
        }
    }

    /// Appends `chunk` to the end of the in-progress `Rope`.
    ///
    /// This method is called repeatedly to incrementally build up a
    /// `Rope`.  The passed text chunk can be as large or small as
    /// desired, but larger chunks are more efficient.
    ///
    /// `chunk` must be valid utf8 text.
    pub fn append(&mut self, chunk: &str) {
        let mut chunk = chunk;

        // Repeatedly chop text off the end of the input, creating
        // leaf nodes out of them and appending them to the tree.
        while !chunk.is_empty() {
            // Get the text for the next leaf
            let (leaf_text, remainder) = self.get_next_leaf_text(chunk);
            chunk = remainder;

            // Append the leaf to the rope
            match leaf_text {
                NextText::None => break,
                NextText::UseBuffer => {
                    let string = NodeText::from_str(&self.buffer);
                    self.append_leaf_node(Arc::new(Node::Leaf(string)));
                    self.buffer.clear();
                }
                NextText::String(s) => {
                    self.append_leaf_node(Arc::new(Node::Leaf(NodeText::from_str(s))));
                }
            }
        }
    }

    /// Finishes the build, and returns the `Rope`.
    ///
    /// Note: this consumes the builder.  If you want to continue building
    /// other ropes with the same prefix, you can clone the builder before
    /// calling `finish()`.
    pub fn finish(mut self) -> Rope {
        // Append the last leaf
        if !self.buffer.is_empty() {
            let string = NodeText::from_str(&self.buffer);
            self.append_leaf_node(Arc::new(Node::Leaf(string)));
        }

        // Zip up all the remaining nodes on the stack
        let mut stack_idx = self.stack.len() - 1;
        while stack_idx >= 1 {
            let node = self.stack.pop().unwrap();
            if let Node::Internal(ref mut children) = *Arc::make_mut(&mut self.stack[stack_idx - 1])
            {
                children.push((node.text_info(), node));
            } else {
                unreachable!();
            }
            stack_idx -= 1;
        }

        // Get root and fix any right-side nodes with too few children.
        let mut root = self.stack.pop().unwrap();
        Arc::make_mut(&mut root).zip_fix_right();

        // Crate the rope, make sure it's well-formed, and return it.
        let mut rope = Rope { root: root };
        rope.pull_up_singular_nodes();
        return rope;
    }

    //-----------------------------------------------------------------

    // Returns (next_leaf_text, remaining_text)
    fn get_next_leaf_text<'a>(&mut self, text: &'a str) -> (NextText<'a>, &'a str) {
        if self.buffer.is_empty() {
            if text.len() > MAX_BYTES {
                // Simplest case: just chop off the end of `text`
                let split_idx = find_good_split_idx(text, MAX_BYTES);
                if split_idx == 0 || split_idx == text.len() {
                    self.buffer.push_str(text);
                    return (NextText::None, "");
                } else {
                    return (NextText::String(&text[..split_idx]), &text[split_idx..]);
                }
            } else {
                self.buffer.push_str(text);
                return (NextText::None, "");
            }
        } else if (text.len() + self.buffer.len()) > MAX_BYTES {
            let split_idx = if self.buffer.len() < MAX_BYTES {
                nearest_internal_grapheme_boundary(text, MAX_BYTES - self.buffer.len())
            } else {
                nearest_internal_grapheme_boundary(text, 0)
            };

            if split_idx == 0 || split_idx == text.len() {
                self.buffer.push_str(text);
                return (NextText::None, "");
            } else {
                self.buffer.push_str(&text[..split_idx]);
                return (NextText::UseBuffer, &text[split_idx..]);
            }
        } else {
            self.buffer.push_str(text);
            return (NextText::None, "");
        }
    }

    fn append_leaf_node(&mut self, leaf: Arc<Node>) {
        let last = self.stack.pop().unwrap();
        match *last {
            Node::Leaf(_) => {
                if last.leaf_text().is_empty() {
                    self.stack.push(leaf);
                } else {
                    let mut children = NodeChildren::new();
                    children.push((last.text_info(), last));
                    children.push((leaf.text_info(), leaf));
                    self.stack.push(Arc::new(Node::Internal(children)));
                }
            }

            Node::Internal(_) => {
                self.stack.push(last);
                let mut left = leaf;
                let mut stack_idx = (self.stack.len() - 1) as isize;
                loop {
                    if stack_idx < 0 {
                        // We're above the root, so do a root split.
                        let mut children = NodeChildren::new();
                        children.push((left.text_info(), left));
                        self.stack.insert(0, Arc::new(Node::Internal(children)));
                        break;
                    } else if self.stack[stack_idx as usize].child_count() < (MAX_CHILDREN - 1) {
                        // There's room to add a child, so do that.
                        Arc::make_mut(&mut self.stack[stack_idx as usize])
                            .children()
                            .push((left.text_info(), left));
                        break;
                    } else {
                        // Not enough room to fit a child, so split.
                        left = Arc::new(Node::Internal(
                            Arc::make_mut(&mut self.stack[stack_idx as usize])
                                .children()
                                .push_split((left.text_info(), left)),
                        ));
                        std::mem::swap(&mut left, &mut self.stack[stack_idx as usize]);
                        stack_idx -= 1;
                    }
                }
            }
        }
    }
}

enum NextText<'a> {
    None,
    UseBuffer,
    String(&'a str),
}

fn find_good_split_idx(text: &str, idx: usize) -> usize {
    if is_grapheme_boundary(text, idx) {
        idx
    } else {
        let prev = prev_grapheme_boundary(text, idx);
        let next = next_grapheme_boundary(text, idx);
        if prev > 0 {
            prev
        } else {
            next
        }
    }
}
<|endoftext|>


<commit_before>extern crate ini;
use self::ini::Ini;
use super::http;

extern crate time;
extern crate fruitbasket;

use std::env;
use std::path;
use std::collections::BTreeMap;

const INIFILE: &'static str = "connectr.ini";
const PORT: u32 = 5432;
const WEB_PORT: u32 = 5676;

#[derive(Default)]
pub struct Settings {
    pub port: u32,
    pub secret: String,
    pub client_id: String,
    pub access_token: Option<String>,
    pub refresh_token: Option<String>,
    pub expire_utc: Option<u64>,
    pub presets: Vec<(String,String)>,
}

fn default_inifile() -> String {
    format!("{}/.{}", env::home_dir().unwrap().display(), INIFILE)
}

fn inifile() -> String {
    // Try to load INI file from home directory
    let path = format!("{}/.{}", env::home_dir().unwrap().display(), INIFILE);
    if path::Path::new(&path).exists() {
        info!("Found config: {}", path);
        return path;
    }

    // Default to looking in current working directory
    let path = INIFILE.to_string();
    if path::Path::new(&path).exists() {
        info!("Found config: {}", path);
        return path;
    }

    String::new()
}

pub fn request_web_config() -> BTreeMap<String,String> {
    let form = format!(r###"
{}
<!DOCTYPE HTML>

<title>Connectr Installation</title>

Connectr Installation

Connectr requires a paid Spotify Premium account and a free Spotify developer application.

If you don't have a Premium account, perhaps try a free trial.

To create your free developer application for Connectr, follow these instructions:

 Go to your Spotify Applications page (login with your Spotify credentials)

 Click "CREATE AN APP" in the upper-right corner

 Enter a name (perhaps "Connectr") and description ("Use Connectr app with my account.")

 Add a Redirect URI: http://127.0.0.1:{}

 Copy your Client ID and Client Secret to the fields below.

 Press the SAVE button at the bottom of Spotify's webpage

 Submit this configuration form

<form method="POST" action="#">

Client ID:	<input type="text" name="client_id">

Client Secret:	<input type="text" name="secret">

Presets:
(optional, one per line)	<textarea rows="10" cols="80"  name="presets" placeholder="First Preset Name=spotify:user:spotify:playlist:37i9dQZEVXboyJ0IJdpcuT"></textarea>

<input type="submit" value="Write config file">

</form>

Config will be saved as: {}

If something goes wrong or changes, edit or delete that file.

"###,
                       "HTTP/1.1 200 OK\r\n\r\n",
                       PORT,
                       default_inifile());
    let reply = format!("{}Configuration saved.  You can close this window.",
                        "HTTP/1.1 200 OK\r\n\r\n");
    let mut config = BTreeMap::<String,String>::new();
    config.insert("port".to_string(), PORT.to_string());
    config.append(&mut http::config_request_local_webserver(WEB_PORT, form, reply));
    config
}

pub fn save_web_config(mut config: BTreeMap<String,String>) -> Ini {
    let mut c = Ini::new();
    let port = config.remove("port").unwrap();
    c.with_section(Some("connectr".to_owned()))
        .set("port", port);
    let secret = config.remove("secret").unwrap_or("<PLACEHOLDER>".to_string());
    let client_id = config.remove("client_id").unwrap_or("<PLACEHOLDER>".to_string());
    let presets = config.remove("presets").unwrap_or(String::new());
    c.with_section(Some("application".to_owned()))
        .set("secret", secret)
        .set("client_id", client_id);
    {
        // TODO: INI uses HashMap, doesn't support maintaining order
        for preset in presets.split("\n") {
            let mut pair = preset.split("=");
            if pair.clone().count() == 2 {
                let key = pair.next().unwrap().trim();
                let value = pair.next().unwrap().trim();
                c.set_to(Some("presets"), key.to_string(), value.to_string());
            }
        }
    }
    c.write_to_file(&default_inifile()).unwrap();
    c
}

pub fn read_settings() -> Option<Settings> {
    info!("Attempting to read config file.");
    let conf = match Ini::load_from_file(&inifile()) {
        Ok(c) => c,
        Err(e) => {
            info!("Load file error: {}", e);
            info!("No config file found.");
            info!("Requesting settings via web form.");
            // Launch a local web server and open a browser to it.  Returns
            // the Spotify configuration.
            let web_config = request_web_config();
            save_web_config(web_config)
        }
    };

    let section = conf.section(Some("connectr".to_owned())).unwrap();
    let port = section.get("port").unwrap().parse().unwrap();

    let section = conf.section(Some("application".to_owned())).unwrap();
    let secret = section.get("secret").unwrap();
    let client_id = section.get("client_id").unwrap();
    if client_id.starts_with('<') || secret.starts_with('<') {
        error!("Invalid or missing configuration.  Cannot continue.");
        info!("");
        info!("ERROR: Spotify Client ID or Secret not set in connectr.ini!");
        info!("");
        info!("Create a Spotify application at https://developer.spotify.com/my-applications/ and");
        info!("add the client ID and secret to connectr.ini.");
        info!("");
        info!("Be sure to add a redirect URI of http://127.0.0.1:<PORT> to your Spotify application,");
        info!("and make sure the port matches in connectr.ini.");
        info!("");
        return None;
    }

    let mut access = None;
    let mut refresh = None;
    let mut expire_utc = None;
    if let Some(section) = conf.section(Some("tokens".to_owned())) {
        access = Some(section.get("access").unwrap().clone());
        refresh = Some(section.get("refresh").unwrap().clone());
        expire_utc = Some(section.get("expire").unwrap().parse().unwrap());
        info!("Read access token from INI!");
    }

    let mut presets = Vec::<(String,String)>::new();
    if let Some(section) = conf.section(Some("presets".to_owned())) {
        for (key, value) in section {
            presets.push((key.to_owned(), value.to_owned()));
        }
    }

    Some(Settings { secret: secret.to_string(), client_id: client_id.to_string(), port: port,
                    access_token: access, refresh_token: refresh, expire_utc: expire_utc,
                    presets: presets})
}

pub type SettingsError = String;
pub fn save_tokens(access: &str, refresh: &str, expire_utc: u64) -> Result<(), SettingsError> {
    let mut conf = Ini::load_from_file(&inifile()).unwrap();
    conf.with_section(Some("tokens".to_owned()))
        .set("access", access)
        .set("refresh", refresh)
        .set("expire", expire_utc.to_string());
    conf.write_to_file(&inifile()).unwrap();
    Ok(())
}
<commit_msg>Accept UTF-8 in web config<commit_after>extern crate ini;
use self::ini::Ini;
use super::http;

extern crate time;
extern crate fruitbasket;

use std::env;
use std::path;
use std::collections::BTreeMap;

const INIFILE: &'static str = "connectr.ini";
const PORT: u32 = 5432;
const WEB_PORT: u32 = 5676;

#[derive(Default)]
pub struct Settings {
    pub port: u32,
    pub secret: String,
    pub client_id: String,
    pub access_token: Option<String>,
    pub refresh_token: Option<String>,
    pub expire_utc: Option<u64>,
    pub presets: Vec<(String,String)>,
}

fn default_inifile() -> String {
    format!("{}/.{}", env::home_dir().unwrap().display(), INIFILE)
}

fn inifile() -> String {
    // Try to load INI file from home directory
    let path = format!("{}/.{}", env::home_dir().unwrap().display(), INIFILE);
    if path::Path::new(&path).exists() {
        info!("Found config: {}", path);
        return path;
    }

    // Default to looking in current working directory
    let path = INIFILE.to_string();
    if path::Path::new(&path).exists() {
        info!("Found config: {}", path);
        return path;
    }

    String::new()
}

pub fn request_web_config() -> BTreeMap<String,String> {
    let form = format!(r###"
{}
<!DOCTYPE HTML>

<title>Connectr Installation</title>

Connectr Installation

Connectr requires a paid Spotify Premium account and a free Spotify developer application.

If you don't have a Premium account, perhaps try a free trial.

To create your free developer application for Connectr, follow these instructions:

 Go to your Spotify Applications page (login with your Spotify credentials)

 Click "CREATE AN APP" in the upper-right corner

 Enter a name (perhaps "Connectr") and description ("Use Connectr app with my account.")

 Add a Redirect URI: http://127.0.0.1:{}

 Copy your Client ID and Client Secret to the fields below.

 Press the SAVE button at the bottom of Spotify's webpage

 Submit this configuration form

<form method="POST" action="#" accept-charset="UTF-8">

Client ID:	<input type="text" name="client_id">

Client Secret:	<input type="text" name="secret">

Presets:
(optional, one per line)	<textarea rows="10" cols="80"  name="presets" placeholder="First Preset Name=spotify:user:spotify:playlist:37i9dQZEVXboyJ0IJdpcuT"></textarea>

<input type="submit" value="Write config file">

</form>

Config will be saved as: {}

If something goes wrong or changes, edit or delete that file.

"###,
                       "HTTP/1.1 200 OK\r\n\r\n",
                       PORT,
                       default_inifile());
    let reply = format!("{}Configuration saved.  You can close this window.",
                        "HTTP/1.1 200 OK\r\n\r\n");
    let mut config = BTreeMap::<String,String>::new();
    config.insert("port".to_string(), PORT.to_string());
    config.append(&mut http::config_request_local_webserver(WEB_PORT, form, reply));
    config
}

pub fn save_web_config(mut config: BTreeMap<String,String>) -> Ini {
    let mut c = Ini::new();
    let port = config.remove("port").unwrap();
    c.with_section(Some("connectr".to_owned()))
        .set("port", port);
    let secret = config.remove("secret").unwrap_or("<PLACEHOLDER>".to_string());
    let client_id = config.remove("client_id").unwrap_or("<PLACEHOLDER>".to_string());
    let presets = config.remove("presets").unwrap_or(String::new());
    c.with_section(Some("application".to_owned()))
        .set("secret", secret)
        .set("client_id", client_id);
    {
        // TODO: INI uses HashMap, doesn't support maintaining order
        for preset in presets.split("\n") {
            let mut pair = preset.split("=");
            if pair.clone().count() == 2 {
                let key = pair.next().unwrap().trim();
                let value = pair.next().unwrap().trim();
                c.set_to(Some("presets"), key.to_string(), value.to_string());
            }
        }
    }
    c.write_to_file(&default_inifile()).unwrap();
    c
}

pub fn read_settings() -> Option<Settings> {
    info!("Attempting to read config file.");
    let conf = match Ini::load_from_file(&inifile()) {
        Ok(c) => c,
        Err(e) => {
            info!("Load file error: {}", e);
            info!("No config file found.");
            info!("Requesting settings via web form.");
            // Launch a local web server and open a browser to it.  Returns
            // the Spotify configuration.
            let web_config = request_web_config();
            save_web_config(web_config)
        }
    };

    let section = conf.section(Some("connectr".to_owned())).unwrap();
    let port = section.get("port").unwrap().parse().unwrap();

    let section = conf.section(Some("application".to_owned())).unwrap();
    let secret = section.get("secret").unwrap();
    let client_id = section.get("client_id").unwrap();
    if client_id.starts_with('<') || secret.starts_with('<') {
        error!("Invalid or missing configuration.  Cannot continue.");
        info!("");
        info!("ERROR: Spotify Client ID or Secret not set in connectr.ini!");
        info!("");
        info!("Create a Spotify application at https://developer.spotify.com/my-applications/ and");
        info!("add the client ID and secret to connectr.ini.");
        info!("");
        info!("Be sure to add a redirect URI of http://127.0.0.1:<PORT> to your Spotify application,");
        info!("and make sure the port matches in connectr.ini.");
        info!("");
        return None;
    }

    let mut access = None;
    let mut refresh = None;
    let mut expire_utc = None;
    if let Some(section) = conf.section(Some("tokens".to_owned())) {
        access = Some(section.get("access").unwrap().clone());
        refresh = Some(section.get("refresh").unwrap().clone());
        expire_utc = Some(section.get("expire").unwrap().parse().unwrap());
        info!("Read access token from INI!");
    }

    let mut presets = Vec::<(String,String)>::new();
    if let Some(section) = conf.section(Some("presets".to_owned())) {
        for (key, value) in section {
            presets.push((key.to_owned(), value.to_owned()));
        }
    }

    Some(Settings { secret: secret.to_string(), client_id: client_id.to_string(), port: port,
                    access_token: access, refresh_token: refresh, expire_utc: expire_utc,
                    presets: presets})
}

pub type SettingsError = String;
pub fn save_tokens(access: &str, refresh: &str, expire_utc: u64) -> Result<(), SettingsError> {
    let mut conf = Ini::load_from_file(&inifile()).unwrap();
    conf.with_section(Some("tokens".to_owned()))
        .set("access", access)
        .set("refresh", refresh)
        .set("expire", expire_utc.to_string());
    conf.write_to_file(&inifile()).unwrap();
    Ok(())
}
<|endoftext|>


<commit_before>use std::net::TcpStream;
use std::sync::mpsc;

use metrics::MetricsRecorder;
use metrics::Timer;
use protocol::cmd::Cmd;
use protocol::cmd::Resp;
use tcp_transport::TcpTransport;

use super::CmdSender;
use super::MetricsSender;
use super::RespReceiver;
use super::RespSender;
use super::TransportId;

pub struct TransportTask {
    id: TransportId,
    cmd_tx: CmdSender,
    met_tx: MetricsSender,
}

impl TransportTask {
    pub fn new(id: TransportId,
               cmd_tx: CmdSender,
               met_tx: MetricsSender)
               -> TransportTask {
        TransportTask {
            id: id,
            cmd_tx: cmd_tx,
            met_tx: met_tx,
        }
    }

    pub fn run(&self, stream: TcpStream) {
        let mut rec = MetricsRecorder::new(self.met_tx.clone());
        let mut transport = TcpTransport::new(stream);
        let (resp_tx, resp_rx): (RespSender, RespReceiver) = mpsc::channel();

        loop {
            // println!("Ready to read command...");
            let rv = {
                let _t = Timer::new(&mut rec, "read_cmd");
                transport.read_cmd()
            };

            // If we couldn't parse the command return an error
            if !rv.is_ok() {
                println!("Failed to read command, returning error");
                transport.write_resp(&Resp::Error);
                return; // Here we just drop the connection
            }

            // Send the command to the driver
            let cmd = rv.unwrap();
            let resp_tx_clone = resp_tx.clone();
            let stats = transport.get_stats_clone();
            {
                let _t = Timer::new(&mut rec, "send_cmd");
                self.cmd_tx
                    .send((self.id, resp_tx_clone, cmd, stats))
                    .unwrap();
            }

            // Obtain a response
            let resp = {
                let _t = Timer::new(&mut rec, "recv_resp");
                resp_rx.recv().unwrap()
            };

            // Return a response
            // println!("Returning response: {:?}", &resp);
            let rv = {
                let _t = Timer::new(&mut rec, "write_resp");
                transport.write_resp(&resp)
            };
            if !rv.is_ok() {
                println!("Failed to write response :(");
            }

            // Now flush metrics outside the request path
            rec.flush_metrics();
        }
    }
}
<commit_msg>use more descriptive metric names<commit_after>use std::net::TcpStream;
use std::sync::mpsc;

use metrics::MetricsRecorder;
use metrics::Timer;
use protocol::cmd::Cmd;
use protocol::cmd::Resp;
use tcp_transport::TcpTransport;

use super::CmdSender;
use super::MetricsSender;
use super::RespReceiver;
use super::RespSender;
use super::TransportId;

pub struct TransportTask {
    id: TransportId,
    cmd_tx: CmdSender,
    met_tx: MetricsSender,
}

impl TransportTask {
    pub fn new(id: TransportId,
               cmd_tx: CmdSender,
               met_tx: MetricsSender)
               -> TransportTask {
        TransportTask {
            id: id,
            cmd_tx: cmd_tx,
            met_tx: met_tx,
        }
    }

    pub fn run(&self, stream: TcpStream) {
        let mut rec = MetricsRecorder::new(self.met_tx.clone());
        let mut transport = TcpTransport::new(stream);
        let (resp_tx, resp_rx): (RespSender, RespReceiver) = mpsc::channel();

        loop {
            // println!("Ready to read command...");
            let rv = {
                let _t = Timer::new(&mut rec, "TransportTask:read_cmd");
                transport.read_cmd()
            };

            // If we couldn't parse the command return an error
            if !rv.is_ok() {
                println!("Failed to read command, returning error");
                transport.write_resp(&Resp::Error);
                return; // Here we just drop the connection
            }

            // Send the command to the driver
            let cmd = rv.unwrap();
            let resp_tx_clone = resp_tx.clone();
            let stats = transport.get_stats_clone();
            {
                let _t = Timer::new(&mut rec, "TransportTask:send_cmd");
                self.cmd_tx
                    .send((self.id, resp_tx_clone, cmd, stats))
                    .unwrap();
            }

            // Obtain a response
            let resp = {
                let _t = Timer::new(&mut rec, "TransportTask:recv_resp");
                resp_rx.recv().unwrap()
            };

            // Return a response
            // println!("Returning response: {:?}", &resp);
            let rv = {
                let _t = Timer::new(&mut rec, "TransportTask:write_resp");
                transport.write_resp(&resp)
            };
            if !rv.is_ok() {
                println!("Failed to write response :(");
            }

            // Now flush metrics outside the request path
            rec.flush_metrics();
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::error::Error;
use std::iter::FromIterator;
use std::time::Duration;

use serialport as sp;
use errors::*;

/// Struct for containing Port information
struct OpenPort {
  /// The opened serial port
  /// SerialPort is not Sized, so it makes hashmap mad
  /// and so we deal with these shennanigans
  port: Box<sp::SerialPort>,
}

impl OpenPort {
  /// Write data to the serial port
  pub fn write_port(&mut self, data: &[u8]) -> Result<()> {
    self
      .port
      .write_all(data)
      .map_err(|err| ErrorKind::Io(err).into())
  }

  /// Read data from the serial port
  pub fn read_port(&mut self, buff: &mut [u8]) -> Result<usize> {
    self
      .port
      .read(buff)
      .map_err(|err| ErrorKind::Io(err).into())
  }
}

/// Manages ports and their locks
pub struct PortManager {
  /// Maintains list of ports
  open_ports: HashMap<String, OpenPort>,
}

impl PortManager {
  /// Create a new PortManager instance
  pub fn new() -> PortManager {
    PortManager { open_ports: HashMap::new() }
  }

  /// Has the port been opened
  pub fn is_port_open(&self, port_name: &String) -> bool {
    self.open_ports.contains_key(port_name)
  }

  /// List all serial ports
  pub fn list_ports(&self) -> Result<Vec<sp::SerialPortInfo>> {
    sp::available_ports().map_err(|e| ErrorKind::Serialport(e).into())
  }

  /// Open a port
  pub fn open_port(&mut self, port_name: &String) -> Result<()> {

    if self.is_port_open(port_name) {
      Ok(())
    } else {

      let sp_settings = sp::SerialPortSettings {
        baud_rate: sp::BaudRate::Baud115200,
        data_bits: sp::DataBits::Eight,
        flow_control: sp::FlowControl::None,
        parity: sp::Parity::None,
        stop_bits: sp::StopBits::One,
        timeout: Duration::from_millis(1),
      };

      match sp::open_with_settings(&port_name, &sp_settings) {
        Ok(serial_port) => {
          let open_port = OpenPort { port: serial_port };
          self.open_ports.insert(port_name.to_string(), open_port);
          Ok(())
        }
        Err(e) => Err(ErrorKind::Serialport(e).into()),
      }
    }
  }

  pub fn close_port(&mut self, port_name: &String) {
    // This drops the underlying serial port and box
    self.open_ports.remove(port_name);
  }

  /// Write data to the port
  pub fn write_port(&mut self, port_name: &String, data: &[u8]) -> Result<()> {
    match self.open_ports.get_mut(port_name) {
      None => Err(ErrorKind::OpenPortNotFound(port_name.to_string()).into()),
      Some(p) => p.write_port(data),
    }
  }

  /// Read data from a port into the buffer buff
  /// If successful, returns Ok(usize) which is the number of
  /// bytes read
  pub fn read_port(&mut self, port_name: &String, buff: &mut [u8]) -> Result<usize> {
    match self.open_ports.get_mut(port_name) {
      None => Err(ErrorKind::OpenPortNotFound(port_name.to_string()).into()),
      Some(p) => p.read_port(buff),
    }
  }

  /// Read all currently open ports, return a hashmap of
  /// ports to Result<Vec<u8>>
  pub fn read_all_ports(&mut self) -> HashMap<String, Result<Vec<u8>>> {
    let mut buffer = vec![0; 4096];
    let mut map = HashMap::new();
    for (port_name, open_port) in self.open_ports.iter_mut() {
      match open_port.read_port(buffer.as_mut_slice()) {
        Ok(bytes_read) => {
          if bytes_read == 0 {
            // EOF
            info!("Received EOF reading from port {}", port_name);
            map.insert(
              port_name.to_string(),
              Err(ErrorKind::PortEOFError(port_name.clone()).into()),
            );
          } else {
            let bytes = buffer[0..bytes_read].to_vec();
            map.insert(port_name.to_string(), Ok(bytes));
          }
        }
        Err(e) => {
          // debug!("Error {} reading from port {}", e, port_name);
          match e.description() {
            "Operation timed out" => {}
            _ => {
              map.insert(port_name.to_string(), Err(e.into()));
            }
          }

        }
      }
    }
    map
  }

  /// Get a vec of open ports
  pub fn open_ports(&self) -> HashSet<String> {
    HashSet::<String>::from_iter(self.open_ports.keys().map(|k| k.clone()))
  }
}

#[cfg(test)]
mod tests {

  use std::io::Write;

  use serialport::SerialPort;
  use serialport::posix::TTYPort;

  use super::*;

  #[test]
  #[cfg(unix)]
  fn test_unix_serialports() {

    let (mut master, mut slave) = TTYPort::pair().expect("Failed to create pseudoterminal pair!");

    slave
      .set_exclusive(false)
      .expect("Failed to set exclusive false");

    let serial_msg = "abcdefg";

    if let Some(s_name) = slave.port_name() {
      let mut port_manager = PortManager::new();

      port_manager
        .open_port(&s_name)
        .expect(&format!("Failed to open slave port {}", s_name));

      master
        .write(serial_msg.as_bytes())
        .expect("Write to master failed!");

      let res = port_manager.read_all_ports();

      assert_eq!(res.len(), 1, "Should have read one port.");

      for (port_name, value) in res {
        match value {
          Ok(bytes) => {
            let read_msg = String::from_utf8_lossy(&bytes);
            assert_eq!(
              serial_msg,
              read_msg,
              "Messages should be same '{}' '{}'",
              serial_msg,
              read_msg
            );
          }
          Err(e) => panic!("Got error reading port {}", e),
        }
      }
    } else {
      panic!("Failed to get slave pty name");
    }
  }

  // TODO: write to slave, read from master

}<commit_msg>Flush after write<commit_after>use std::collections::{HashMap, HashSet};
use std::error::Error;
use std::iter::FromIterator;
use std::time::Duration;

use serialport as sp;
use errors::*;

/// Struct for containing Port information
struct OpenPort {
  /// The opened serial port
  /// SerialPort is not Sized, so it makes hashmap mad
  /// and so we deal with these shennanigans
  port: Box<sp::SerialPort>,
}

impl OpenPort {
  /// Write data to the serial port
  pub fn write_port(&mut self, data: &[u8]) -> Result<()> {
    self
      .port
      .write_all(data)
      .and_then(|_| self.port.flush())
      .map_err(|err| ErrorKind::Io(err).into())
  }

  /// Read data from the serial port
  pub fn read_port(&mut self, buff: &mut [u8]) -> Result<usize> {
    self
      .port
      .read(buff)
      .map_err(|err| ErrorKind::Io(err).into())
  }
}

/// Manages ports and their locks
pub struct PortManager {
  /// Maintains list of ports
  open_ports: HashMap<String, OpenPort>,
}

impl PortManager {
  /// Create a new PortManager instance
  pub fn new() -> PortManager {
    PortManager { open_ports: HashMap::new() }
  }

  /// Has the port been opened
  pub fn is_port_open(&self, port_name: &String) -> bool {
    self.open_ports.contains_key(port_name)
  }

  /// List all serial ports
  pub fn list_ports(&self) -> Result<Vec<sp::SerialPortInfo>> {
    sp::available_ports().map_err(|e| ErrorKind::Serialport(e).into())
  }

  /// Open a port
  pub fn open_port(&mut self, port_name: &String) -> Result<()> {

    if self.is_port_open(port_name) {
      Ok(())
    } else {

      let sp_settings = sp::SerialPortSettings {
        baud_rate: sp::BaudRate::Baud115200,
        data_bits: sp::DataBits::Eight,
        flow_control: sp::FlowControl::None,
        parity: sp::Parity::None,
        stop_bits: sp::StopBits::One,
        timeout: Duration::from_millis(1),
      };

      match sp::open_with_settings(&port_name, &sp_settings) {
        Ok(serial_port) => {
          let open_port = OpenPort { port: serial_port };
          self.open_ports.insert(port_name.to_string(), open_port);
          Ok(())
        }
        Err(e) => Err(ErrorKind::Serialport(e).into()),
      }
    }
  }

  pub fn close_port(&mut self, port_name: &String) {
    // This drops the underlying serial port and box
    self.open_ports.remove(port_name);
  }

  /// Write data to the port
  pub fn write_port(&mut self, port_name: &String, data: &[u8]) -> Result<()> {
    match self.open_ports.get_mut(port_name) {
      None => Err(ErrorKind::OpenPortNotFound(port_name.to_string()).into()),
      Some(p) => p.write_port(data),
    }
  }

  /// Read data from a port into the buffer buff
  /// If successful, returns Ok(usize) which is the number of
  /// bytes read
  pub fn read_port(&mut self, port_name: &String, buff: &mut [u8]) -> Result<usize> {
    match self.open_ports.get_mut(port_name) {
      None => Err(ErrorKind::OpenPortNotFound(port_name.to_string()).into()),
      Some(p) => p.read_port(buff),
    }
  }

  /// Read all currently open ports, return a hashmap of
  /// ports to Result<Vec<u8>>
  pub fn read_all_ports(&mut self) -> HashMap<String, Result<Vec<u8>>> {
    let mut buffer = vec![0; 4096];
    let mut map = HashMap::new();
    for (port_name, open_port) in self.open_ports.iter_mut() {
      match open_port.read_port(buffer.as_mut_slice()) {
        Ok(bytes_read) => {
          if bytes_read == 0 {
            // EOF
            info!("Received EOF reading from port {}", port_name);
            map.insert(
              port_name.to_string(),
              Err(ErrorKind::PortEOFError(port_name.clone()).into()),
            );
          } else {
            let bytes = buffer[0..bytes_read].to_vec();
            map.insert(port_name.to_string(), Ok(bytes));
          }
        }
        Err(e) => {
          // debug!("Error {} reading from port {}", e, port_name);
          match e.description() {
            "Operation timed out" => {}
            _ => {
              map.insert(port_name.to_string(), Err(e.into()));
            }
          }

        }
      }
    }
    map
  }

  /// Get a vec of open ports
  pub fn open_ports(&self) -> HashSet<String> {
    HashSet::<String>::from_iter(self.open_ports.keys().map(|k| k.clone()))
  }
}

#[cfg(test)]
mod tests {

  use std::io::Write;

  use serialport::SerialPort;
  use serialport::posix::TTYPort;

  use super::*;

  #[test]
  #[cfg(unix)]
  fn test_unix_serialports() {

    let (mut master, mut slave) = TTYPort::pair().expect("Failed to create pseudoterminal pair!");

    slave
      .set_exclusive(false)
      .expect("Failed to set exclusive false");

    let serial_msg = "abcdefg";

    if let Some(s_name) = slave.port_name() {
      let mut port_manager = PortManager::new();

      port_manager
        .open_port(&s_name)
        .expect(&format!("Failed to open slave port {}", s_name));

      master
        .write(serial_msg.as_bytes())
        .expect("Write to master failed!");

      let res = port_manager.read_all_ports();

      assert_eq!(res.len(), 1, "Should have read one port.");

      for (port_name, value) in res {
        match value {
          Ok(bytes) => {
            let read_msg = String::from_utf8_lossy(&bytes);
            assert_eq!(
              serial_msg,
              read_msg,
              "Messages should be same '{}' '{}'",
              serial_msg,
              read_msg
            );
          }
          Err(e) => panic!("Got error reading port {}", e),
        }
      }
    } else {
      panic!("Failed to get slave pty name");
    }
  }

  // TODO: write to slave, read from master

}<|endoftext|>


<commit_before>use std::iter;

use rustc_middle::mir;
use rustc_target::abi::Size;

use crate::*;

impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn emulate_foreign_item_by_name(
        &mut self,
        link_name: &str,
        args: &[OpTy<'tcx, Tag>],
        dest: PlaceTy<'tcx, Tag>,
        _ret: mir::BasicBlock,
    ) -> InterpResult<'tcx, bool> {
        let this = self.eval_context_mut();

        // Windows API stubs.
        // HANDLE = isize
        // DWORD = ULONG = u32
        // BOOL = i32
        match link_name {
            // Environment related shims
            "GetEnvironmentVariableW" => {
                let result = this.GetEnvironmentVariableW(args[0], args[1], args[2])?;
                this.write_scalar(Scalar::from_u32(result), dest)?;
            }
            "SetEnvironmentVariableW" => {
                let result = this.SetEnvironmentVariableW(args[0], args[1])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }
            "GetEnvironmentStringsW" => {
                let result = this.GetEnvironmentStringsW()?;
                this.write_scalar(result, dest)?;
            }
            "FreeEnvironmentStringsW" => {
                let result = this.FreeEnvironmentStringsW(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }
            "GetCurrentDirectoryW" => {
                let result = this.GetCurrentDirectoryW(args[0], args[1])?;
                this.write_scalar(Scalar::from_u32(result), dest)?;
            }
            "SetCurrentDirectoryW" => {
                let result = this.SetCurrentDirectoryW(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }

            // File related shims
            "GetStdHandle" => {
                let which = this.read_scalar(args[0])?.to_i32()?;
                // We just make this the identity function, so we know later in `WriteFile`
                // which one it is.
                this.write_scalar(Scalar::from_machine_isize(which.into(), this), dest)?;
            }
            "WriteFile" => {
                let handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let buf = this.read_scalar(args[1])?.not_undef()?;
                let n = this.read_scalar(args[2])?.to_u32()?;
                let written_place = this.deref_operand(args[3])?;
                // Spec says to always write `0` first.
                this.write_null(written_place.into())?;
                let written = if handle == -11 || handle == -12 {
                    // stdout/stderr
                    use std::io::{self, Write};

                    let buf_cont = this.memory.read_bytes(buf, Size::from_bytes(u64::from(n)))?;
                    let res = if handle == -11 {
                        io::stdout().write(buf_cont)
                    } else {
                        io::stderr().write(buf_cont)
                    };
                    res.ok().map(|n| n as u32)
                } else {
                    throw_unsup_format!("on Windows, writing to anything except stdout/stderr is not supported")
                };
                // If there was no error, write back how much was written.
                if let Some(n) = written {
                    this.write_scalar(Scalar::from_u32(n), written_place.into())?;
                }
                // Return whether this was a success.
                this.write_scalar(
                    Scalar::from_i32(if written.is_some() { 1 } else { 0 }),
                    dest,
                )?;
            }

            // Allocation
            "HeapAlloc" => {
                let _handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let flags = this.read_scalar(args[1])?.to_u32()?;
                let size = this.read_scalar(args[2])?.to_machine_usize(this)?;
                let zero_init = (flags & 0x00000008) != 0; // HEAP_ZERO_MEMORY
                let res = this.malloc(size, zero_init, MiriMemoryKind::WinHeap);
                this.write_scalar(res, dest)?;
            }
            "HeapFree" => {
                let _handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _flags = this.read_scalar(args[1])?.to_u32()?;
                let ptr = this.read_scalar(args[2])?.not_undef()?;
                this.free(ptr, MiriMemoryKind::WinHeap)?;
                this.write_scalar(Scalar::from_i32(1), dest)?;
            }
            "HeapReAlloc" => {
                let _handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _flags = this.read_scalar(args[1])?.to_u32()?;
                let ptr = this.read_scalar(args[2])?.not_undef()?;
                let size = this.read_scalar(args[3])?.to_machine_usize(this)?;
                let res = this.realloc(ptr, size, MiriMemoryKind::WinHeap)?;
                this.write_scalar(res, dest)?;
            }

            // errno
            "SetLastError" => {
                this.set_last_error(this.read_scalar(args[0])?.not_undef()?)?;
            }
            "GetLastError" => {
                let last_error = this.get_last_error()?;
                this.write_scalar(last_error, dest)?;
            }

            // Querying system information
            "GetSystemInfo" => {
                let system_info = this.deref_operand(args[0])?;
                // Initialize with `0`.
                this.memory.write_bytes(
                    system_info.ptr,
                    iter::repeat(0u8).take(system_info.layout.size.bytes() as usize),
                )?;
                // Set number of processors.
                let dword_size = Size::from_bytes(4);
                let num_cpus = this.mplace_field(system_info, 6)?;
                this.write_scalar(Scalar::from_int(NUM_CPUS, dword_size), num_cpus.into())?;
            }

            // Thread-local storage
            "TlsAlloc" => {
                // This just creates a key; Windows does not natively support TLS destructors.

                // Create key and return it.
                let key = this.machine.tls.create_tls_key(None, dest.layout.size)?;
                this.write_scalar(Scalar::from_uint(key, dest.layout.size), dest)?;
            }
            "TlsGetValue" => {
                let key = u128::from(this.read_scalar(args[0])?.to_u32()?);
                let active_thread = this.get_active_thread()?;
                let ptr = this.machine.tls.load_tls(key, active_thread, this)?;
                this.write_scalar(ptr, dest)?;
            }
            "TlsSetValue" => {
                let key = u128::from(this.read_scalar(args[0])?.to_u32()?);
                let active_thread = this.get_active_thread()?;
                let new_ptr = this.read_scalar(args[1])?.not_undef()?;
                this.machine.tls.store_tls(key, active_thread, this.test_null(new_ptr)?)?;

                // Return success (`1`).
                this.write_scalar(Scalar::from_i32(1), dest)?;
            }

            // Access to command-line arguments
            "GetCommandLineW" => {
                this.write_scalar(
                    this.machine.cmd_line.expect("machine must be initialized"),
                    dest,
                )?;
            }

            // Time related shims
            "GetSystemTimeAsFileTime" => {
                this.GetSystemTimeAsFileTime(args[0])?;
            }
            "QueryPerformanceCounter" => {
                let result = this.QueryPerformanceCounter(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }
            "QueryPerformanceFrequency" => {
                let result = this.QueryPerformanceFrequency(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }

            // Miscellaneous
            "SystemFunction036" => {
                // The actual name of 'RtlGenRandom'
                let ptr = this.read_scalar(args[0])?.not_undef()?;
                let len = this.read_scalar(args[1])?.to_u32()?;
                this.gen_random(ptr, len.into())?;
                this.write_scalar(Scalar::from_bool(true), dest)?;
            }
            "GetConsoleScreenBufferInfo" => {
                // `term` needs this, so we fake it.
                let _console = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _buffer_info = this.deref_operand(args[1])?;
                // Indicate an error.
                // FIXME: we should set last_error, but to what?
                this.write_null(dest)?;
            }
            "GetConsoleMode" => {
                // Windows "isatty" (in libtest) needs this, so we fake it.
                let _console = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _mode = this.deref_operand(args[1])?;
                // Indicate an error.
                // FIXME: we should set last_error, but to what?
                this.write_null(dest)?;
            }
            "SwitchToThread" => {
                // Note that once Miri supports concurrency, this will need to return a nonzero
                // value if this call does result in switching to another thread.
                this.write_null(dest)?;
            }

            // Better error for attempts to create a thread
            "CreateThread" => {
                throw_unsup_format!("Miri does not support threading");
            }

            // Incomplete shims that we "stub out" just to get pre-main initialization code to work.
            // These shims are enabled only when the caller is in the standard library.
            "GetProcessHeap" if this.frame().instance.to_string().starts_with("std::sys::windows::") => {
                // Just fake a HANDLE
                this.write_scalar(Scalar::from_machine_isize(1, this), dest)?;
            }
            | "GetModuleHandleW"
            | "GetProcAddress"
            | "SetConsoleTextAttribute" if this.frame().instance.to_string().starts_with("std::sys::windows::")
            => {
                // Pretend these do not exist / nothing happened, by returning zero.
                this.write_null(dest)?;
            }
            "AddVectoredExceptionHandler" if this.frame().instance.to_string().starts_with("std::sys::windows::") => {
                // Any non zero value works for the stdlib. This is just used for stack overflows anyway.
                this.write_scalar(Scalar::from_machine_usize(1, this), dest)?;
            }
            | "InitializeCriticalSection"
            | "EnterCriticalSection"
            | "LeaveCriticalSection"
            | "DeleteCriticalSection" if this.frame().instance.to_string().starts_with("std::sys::windows::")
            => {
                // Nothing to do, not even a return value.
                // (Windows locks are reentrant, and we have only 1 thread,
                // so not doing any futher checks here is at least not incorrect.)
            }
            "TryEnterCriticalSection" if this.frame().instance.to_string().starts_with("std::sys::windows::")
            => {
                // There is only one thread, so this always succeeds and returns TRUE
                this.write_scalar(Scalar::from_i32(1), dest)?;
            }

            _ => throw_unsup_format!("can't call foreign function: {}", link_name),
        }

        Ok(true)
    }
}

<commit_msg>Windows lock primitives: check that we are truly sinle-threaded<commit_after>use std::iter;

use rustc_middle::mir;
use rustc_target::abi::Size;

use crate::*;

impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn emulate_foreign_item_by_name(
        &mut self,
        link_name: &str,
        args: &[OpTy<'tcx, Tag>],
        dest: PlaceTy<'tcx, Tag>,
        _ret: mir::BasicBlock,
    ) -> InterpResult<'tcx, bool> {
        let this = self.eval_context_mut();

        // Windows API stubs.
        // HANDLE = isize
        // DWORD = ULONG = u32
        // BOOL = i32
        match link_name {
            // Environment related shims
            "GetEnvironmentVariableW" => {
                let result = this.GetEnvironmentVariableW(args[0], args[1], args[2])?;
                this.write_scalar(Scalar::from_u32(result), dest)?;
            }
            "SetEnvironmentVariableW" => {
                let result = this.SetEnvironmentVariableW(args[0], args[1])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }
            "GetEnvironmentStringsW" => {
                let result = this.GetEnvironmentStringsW()?;
                this.write_scalar(result, dest)?;
            }
            "FreeEnvironmentStringsW" => {
                let result = this.FreeEnvironmentStringsW(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }
            "GetCurrentDirectoryW" => {
                let result = this.GetCurrentDirectoryW(args[0], args[1])?;
                this.write_scalar(Scalar::from_u32(result), dest)?;
            }
            "SetCurrentDirectoryW" => {
                let result = this.SetCurrentDirectoryW(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }

            // File related shims
            "GetStdHandle" => {
                let which = this.read_scalar(args[0])?.to_i32()?;
                // We just make this the identity function, so we know later in `WriteFile`
                // which one it is.
                this.write_scalar(Scalar::from_machine_isize(which.into(), this), dest)?;
            }
            "WriteFile" => {
                let handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let buf = this.read_scalar(args[1])?.not_undef()?;
                let n = this.read_scalar(args[2])?.to_u32()?;
                let written_place = this.deref_operand(args[3])?;
                // Spec says to always write `0` first.
                this.write_null(written_place.into())?;
                let written = if handle == -11 || handle == -12 {
                    // stdout/stderr
                    use std::io::{self, Write};

                    let buf_cont = this.memory.read_bytes(buf, Size::from_bytes(u64::from(n)))?;
                    let res = if handle == -11 {
                        io::stdout().write(buf_cont)
                    } else {
                        io::stderr().write(buf_cont)
                    };
                    res.ok().map(|n| n as u32)
                } else {
                    throw_unsup_format!("on Windows, writing to anything except stdout/stderr is not supported")
                };
                // If there was no error, write back how much was written.
                if let Some(n) = written {
                    this.write_scalar(Scalar::from_u32(n), written_place.into())?;
                }
                // Return whether this was a success.
                this.write_scalar(
                    Scalar::from_i32(if written.is_some() { 1 } else { 0 }),
                    dest,
                )?;
            }

            // Allocation
            "HeapAlloc" => {
                let _handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let flags = this.read_scalar(args[1])?.to_u32()?;
                let size = this.read_scalar(args[2])?.to_machine_usize(this)?;
                let zero_init = (flags & 0x00000008) != 0; // HEAP_ZERO_MEMORY
                let res = this.malloc(size, zero_init, MiriMemoryKind::WinHeap);
                this.write_scalar(res, dest)?;
            }
            "HeapFree" => {
                let _handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _flags = this.read_scalar(args[1])?.to_u32()?;
                let ptr = this.read_scalar(args[2])?.not_undef()?;
                this.free(ptr, MiriMemoryKind::WinHeap)?;
                this.write_scalar(Scalar::from_i32(1), dest)?;
            }
            "HeapReAlloc" => {
                let _handle = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _flags = this.read_scalar(args[1])?.to_u32()?;
                let ptr = this.read_scalar(args[2])?.not_undef()?;
                let size = this.read_scalar(args[3])?.to_machine_usize(this)?;
                let res = this.realloc(ptr, size, MiriMemoryKind::WinHeap)?;
                this.write_scalar(res, dest)?;
            }

            // errno
            "SetLastError" => {
                this.set_last_error(this.read_scalar(args[0])?.not_undef()?)?;
            }
            "GetLastError" => {
                let last_error = this.get_last_error()?;
                this.write_scalar(last_error, dest)?;
            }

            // Querying system information
            "GetSystemInfo" => {
                let system_info = this.deref_operand(args[0])?;
                // Initialize with `0`.
                this.memory.write_bytes(
                    system_info.ptr,
                    iter::repeat(0u8).take(system_info.layout.size.bytes() as usize),
                )?;
                // Set number of processors.
                let dword_size = Size::from_bytes(4);
                let num_cpus = this.mplace_field(system_info, 6)?;
                this.write_scalar(Scalar::from_int(NUM_CPUS, dword_size), num_cpus.into())?;
            }

            // Thread-local storage
            "TlsAlloc" => {
                // This just creates a key; Windows does not natively support TLS destructors.

                // Create key and return it.
                let key = this.machine.tls.create_tls_key(None, dest.layout.size)?;
                this.write_scalar(Scalar::from_uint(key, dest.layout.size), dest)?;
            }
            "TlsGetValue" => {
                let key = u128::from(this.read_scalar(args[0])?.to_u32()?);
                let active_thread = this.get_active_thread()?;
                let ptr = this.machine.tls.load_tls(key, active_thread, this)?;
                this.write_scalar(ptr, dest)?;
            }
            "TlsSetValue" => {
                let key = u128::from(this.read_scalar(args[0])?.to_u32()?);
                let active_thread = this.get_active_thread()?;
                let new_ptr = this.read_scalar(args[1])?.not_undef()?;
                this.machine.tls.store_tls(key, active_thread, this.test_null(new_ptr)?)?;

                // Return success (`1`).
                this.write_scalar(Scalar::from_i32(1), dest)?;
            }

            // Access to command-line arguments
            "GetCommandLineW" => {
                this.write_scalar(
                    this.machine.cmd_line.expect("machine must be initialized"),
                    dest,
                )?;
            }

            // Time related shims
            "GetSystemTimeAsFileTime" => {
                this.GetSystemTimeAsFileTime(args[0])?;
            }
            "QueryPerformanceCounter" => {
                let result = this.QueryPerformanceCounter(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }
            "QueryPerformanceFrequency" => {
                let result = this.QueryPerformanceFrequency(args[0])?;
                this.write_scalar(Scalar::from_i32(result), dest)?;
            }

            // Miscellaneous
            "SystemFunction036" => {
                // The actual name of 'RtlGenRandom'
                let ptr = this.read_scalar(args[0])?.not_undef()?;
                let len = this.read_scalar(args[1])?.to_u32()?;
                this.gen_random(ptr, len.into())?;
                this.write_scalar(Scalar::from_bool(true), dest)?;
            }
            "GetConsoleScreenBufferInfo" => {
                // `term` needs this, so we fake it.
                let _console = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _buffer_info = this.deref_operand(args[1])?;
                // Indicate an error.
                // FIXME: we should set last_error, but to what?
                this.write_null(dest)?;
            }
            "GetConsoleMode" => {
                // Windows "isatty" (in libtest) needs this, so we fake it.
                let _console = this.read_scalar(args[0])?.to_machine_isize(this)?;
                let _mode = this.deref_operand(args[1])?;
                // Indicate an error.
                // FIXME: we should set last_error, but to what?
                this.write_null(dest)?;
            }
            "SwitchToThread" => {
                // Note that once Miri supports concurrency, this will need to return a nonzero
                // value if this call does result in switching to another thread.
                this.write_null(dest)?;
            }

            // Better error for attempts to create a thread
            "CreateThread" => {
                throw_unsup_format!("Miri does not support threading");
            }

            // Incomplete shims that we "stub out" just to get pre-main initialization code to work.
            // These shims are enabled only when the caller is in the standard library.
            "GetProcessHeap" if this.frame().instance.to_string().starts_with("std::sys::windows::") => {
                // Just fake a HANDLE
                this.write_scalar(Scalar::from_machine_isize(1, this), dest)?;
            }
            | "GetModuleHandleW"
            | "GetProcAddress"
            | "SetConsoleTextAttribute" if this.frame().instance.to_string().starts_with("std::sys::windows::")
            => {
                // Pretend these do not exist / nothing happened, by returning zero.
                this.write_null(dest)?;
            }
            "AddVectoredExceptionHandler" if this.frame().instance.to_string().starts_with("std::sys::windows::") => {
                // Any non zero value works for the stdlib. This is just used for stack overflows anyway.
                this.write_scalar(Scalar::from_machine_usize(1, this), dest)?;
            }
            | "InitializeCriticalSection"
            | "EnterCriticalSection"
            | "LeaveCriticalSection"
            | "DeleteCriticalSection" if this.frame().instance.to_string().starts_with("std::sys::windows::")
            => {
                assert_eq!(this.get_total_thread_count()?, 1, "concurrency on Windows not supported");
                // Nothing to do, not even a return value.
                // (Windows locks are reentrant, and we have only 1 thread,
                // so not doing any futher checks here is at least not incorrect.)
            }
            "TryEnterCriticalSection" if this.frame().instance.to_string().starts_with("std::sys::windows::")
            => {
                assert_eq!(this.get_total_thread_count()?, 1, "concurrency on Windows not supported");
                // There is only one thread, so this always succeeds and returns TRUE
                this.write_scalar(Scalar::from_i32(1), dest)?;
            }

            _ => throw_unsup_format!("can't call foreign function: {}", link_name),
        }

        Ok(true)
    }
}

<|endoftext|>


<commit_before>// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.

use std::collections::hash_map::DefaultHasher;
use std::collections::VecDeque;
use std::hash::{Hash, Hasher};
use std::sync::Mutex;
use std::usize;

/// Latch which is used to serialize accesses to resources hashed to the same slot.
///
/// Latches are indexed by slot IDs. The keys of a command are hashed to slot IDs, then the command
/// is added to the waiting queues of the latches.
///
/// If command A is ahead of command B in one latch, it must be ahead of command B in all the
/// overlapping latches. This is an invariant ensured by the `gen_lock`, `acquire` and `release`.
#[derive(Clone)]
struct Latch {
    // store waiting commands
    pub waiting: VecDeque<u64>,
}

impl Latch {
    /// Creates a latch with an empty waiting queue.
    pub fn new() -> Latch {
        Latch {
            waiting: VecDeque::new(),
        }
    }
}

/// Lock required for a command.
#[derive(Clone)]
pub struct Lock {
    /// The slot IDs of the latches that a command must acquire before being able to be processed.
    pub required_slots: Vec<usize>,

    /// The number of latches that the command has acquired.
    pub owned_count: usize,
}

impl Lock {
    /// Creates a lock.
    pub fn new(required_slots: Vec<usize>) -> Lock {
        Lock {
            required_slots,
            owned_count: 0,
        }
    }

    /// Returns true if all the required latches have be acquired, false otherwise.
    pub fn acquired(&self) -> bool {
        self.required_slots.len() == self.owned_count
    }

    pub fn is_write_lock(&self) -> bool {
        !self.required_slots.is_empty()
    }
}

/// Latches which are used for concurrency control in the scheduler.
///
/// Each latch is indexed by a slot ID, hence the term latch and slot are used interchangeably, but
/// conceptually a latch is a queue, and a slot is an index to the queue.
pub struct Latches {
    slots: Vec<Mutex<Latch>>,
    size: usize,
}

impl Latches {
    /// Creates latches.
    ///
    /// The size will be rounded up to the power of 2.
    pub fn new(size: usize) -> Latches {
        let size = usize::next_power_of_two(size);
        let mut slots = Vec::with_capacity(size);
        (0..size).for_each(|_| slots.push(Mutex::new(Latch::new())));
        Latches { slots, size }
    }

    /// Creates a lock which specifies all the required latches for a command.
    pub fn gen_lock<H>(&self, keys: &[H]) -> Lock
    where
        H: Hash,
    {
        // prevent from deadlock, so we sort and deduplicate the index
        let mut slots: Vec<usize> = keys.iter().map(|x| self.calc_slot(x)).collect();
        slots.sort();
        slots.dedup();
        Lock::new(slots)
    }

    /// Tries to acquire the latches specified by the `lock` for command with ID `who`.
    ///
    /// This method will enqueue the command ID into the waiting queues of the latches. A latch is
    /// considered acquired if the command ID is at the front of the queue. Returns true if all the
    /// Latches are acquired, false otherwise.
    pub fn acquire(&self, lock: &mut Lock, who: u64) -> bool {
        let mut acquired_count: usize = 0;
        for i in &lock.required_slots[lock.owned_count..] {
            let mut latch = self.slots[*i].lock().unwrap();
            let front = latch.waiting.front().cloned();
            match front {
                Some(cid) => {
                    if cid == who {
                        acquired_count += 1;
                    } else {
                        latch.waiting.push_back(who);
                        break;
                    }
                }
                None => {
                    latch.waiting.push_back(who);
                    acquired_count += 1;
                }
            }
        }
        lock.owned_count += acquired_count;
        lock.acquired()
    }

    /// Releases all latches owned by the `lock` of command with ID `who`, returns the wakeup list.
    ///
    /// Preconditions: the caller must ensure the command is at the front of the latches.
    pub fn release(&self, lock: &Lock, who: u64) -> Vec<u64> {
        let mut wakeup_list: Vec<u64> = vec![];
        for i in &lock.required_slots[..lock.owned_count] {
            let mut latch = self.slots[*i].lock().unwrap();
            let front = latch.waiting.pop_front().unwrap();
            assert_eq!(front, who);
            if let Some(wakeup) = latch.waiting.front() {
                wakeup_list.push(*wakeup);
            }
        }
        wakeup_list
    }

    /// Calculates the slot ID by hashing the `key`.
    fn calc_slot<H>(&self, key: &H) -> usize
    where
        H: Hash,
    {
        let mut s = DefaultHasher::new();
        key.hash(&mut s);
        (s.finish() as usize) & (self.size - 1)
    }
}

#[cfg(test)]
mod tests {
    use super::{Latches, Lock};

    #[test]
    fn test_wakeup() {
        let latches = Latches::new(256);

        let slots_a: Vec<usize> = vec![1, 3, 5];
        let mut lock_a = Lock::new(slots_a);
        let slots_b: Vec<usize> = vec![4, 5, 6];
        let mut lock_b = Lock::new(slots_b);
        let cid_a: u64 = 1;
        let cid_b: u64 = 2;

        // a acquire lock success
        let acquired_a = latches.acquire(&mut lock_a, cid_a);
        assert_eq!(acquired_a, true);

        // b acquire lock failed
        let mut acquired_b = latches.acquire(&mut lock_b, cid_b);
        assert_eq!(acquired_b, false);

        // a release lock, and get wakeup list
        let wakeup = latches.release(&lock_a, cid_a);
        assert_eq!(wakeup[0], cid_b);

        // b acquire lock success
        acquired_b = latches.acquire(&mut lock_b, cid_b);
        assert_eq!(acquired_b, true);
    }

    #[test]
    fn test_wakeup_by_multi_cmds() {
        let latches = Latches::new(256);

        let slots_a: Vec<usize> = vec![1, 2, 3];
        let slots_b: Vec<usize> = vec![4, 5, 6];
        let slots_c: Vec<usize> = vec![3, 4];
        let mut lock_a = Lock::new(slots_a);
        let mut lock_b = Lock::new(slots_b);
        let mut lock_c = Lock::new(slots_c);
        let cid_a: u64 = 1;
        let cid_b: u64 = 2;
        let cid_c: u64 = 3;

        // a acquire lock success
        let acquired_a = latches.acquire(&mut lock_a, cid_a);
        assert_eq!(acquired_a, true);

        // b acquire lock success
        let acquired_b = latches.acquire(&mut lock_b, cid_b);
        assert_eq!(acquired_b, true);

        // c acquire lock failed, cause a occupied slot 3
        let mut acquired_c = latches.acquire(&mut lock_c, cid_c);
        assert_eq!(acquired_c, false);

        // a release lock, and get wakeup list
        let wakeup = latches.release(&lock_a, cid_a);
        assert_eq!(wakeup[0], cid_c);

        // c acquire lock failed again, cause b occupied slot 4
        acquired_c = latches.acquire(&mut lock_c, cid_c);
        assert_eq!(acquired_c, false);

        // b release lock, and get wakeup list
        let wakeup = latches.release(&lock_b, cid_b);
        assert_eq!(wakeup[0], cid_c);

        // finally c acquire lock success
        acquired_c = latches.acquire(&mut lock_c, cid_c);
        assert_eq!(acquired_c, true);
    }
}
<commit_msg>shrink the latch waiting list (#4844)<commit_after>// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.

use std::collections::hash_map::DefaultHasher;
use std::collections::VecDeque;
use std::hash::{Hash, Hasher};
use std::sync::Mutex;
use std::usize;

const WAITING_LIST_SHRINK_SIZE: usize = 32;

/// Latch which is used to serialize accesses to resources hashed to the same slot.
///
/// Latches are indexed by slot IDs. The keys of a command are hashed to slot IDs, then the command
/// is added to the waiting queues of the latches.
///
/// If command A is ahead of command B in one latch, it must be ahead of command B in all the
/// overlapping latches. This is an invariant ensured by the `gen_lock`, `acquire` and `release`.
#[derive(Clone)]
struct Latch {
    // store waiting commands
    pub waiting: VecDeque<u64>,
}

impl Latch {
    /// Creates a latch with an empty waiting queue.
    pub fn new() -> Latch {
        Latch {
            waiting: VecDeque::new(),
        }
    }
}

/// Lock required for a command.
#[derive(Clone)]
pub struct Lock {
    /// The slot IDs of the latches that a command must acquire before being able to be processed.
    pub required_slots: Vec<usize>,

    /// The number of latches that the command has acquired.
    pub owned_count: usize,
}

impl Lock {
    /// Creates a lock.
    pub fn new(required_slots: Vec<usize>) -> Lock {
        Lock {
            required_slots,
            owned_count: 0,
        }
    }

    /// Returns true if all the required latches have be acquired, false otherwise.
    pub fn acquired(&self) -> bool {
        self.required_slots.len() == self.owned_count
    }

    pub fn is_write_lock(&self) -> bool {
        !self.required_slots.is_empty()
    }
}

/// Latches which are used for concurrency control in the scheduler.
///
/// Each latch is indexed by a slot ID, hence the term latch and slot are used interchangeably, but
/// conceptually a latch is a queue, and a slot is an index to the queue.
pub struct Latches {
    slots: Vec<Mutex<Latch>>,
    size: usize,
}

impl Latches {
    /// Creates latches.
    ///
    /// The size will be rounded up to the power of 2.
    pub fn new(size: usize) -> Latches {
        let size = usize::next_power_of_two(size);
        let mut slots = Vec::with_capacity(size);
        (0..size).for_each(|_| slots.push(Mutex::new(Latch::new())));
        Latches { slots, size }
    }

    /// Creates a lock which specifies all the required latches for a command.
    pub fn gen_lock<H>(&self, keys: &[H]) -> Lock
    where
        H: Hash,
    {
        // prevent from deadlock, so we sort and deduplicate the index
        let mut slots: Vec<usize> = keys.iter().map(|x| self.calc_slot(x)).collect();
        slots.sort();
        slots.dedup();
        Lock::new(slots)
    }

    /// Tries to acquire the latches specified by the `lock` for command with ID `who`.
    ///
    /// This method will enqueue the command ID into the waiting queues of the latches. A latch is
    /// considered acquired if the command ID is at the front of the queue. Returns true if all the
    /// Latches are acquired, false otherwise.
    pub fn acquire(&self, lock: &mut Lock, who: u64) -> bool {
        let mut acquired_count: usize = 0;
        for i in &lock.required_slots[lock.owned_count..] {
            let mut latch = self.slots[*i].lock().unwrap();
            let front = latch.waiting.front().cloned();
            match front {
                Some(cid) => {
                    if cid == who {
                        acquired_count += 1;
                    } else {
                        latch.waiting.push_back(who);
                        break;
                    }
                }
                None => {
                    latch.waiting.push_back(who);
                    acquired_count += 1;
                }
            }
        }
        lock.owned_count += acquired_count;
        lock.acquired()
    }

    /// Releases all latches owned by the `lock` of command with ID `who`, returns the wakeup list.
    ///
    /// Preconditions: the caller must ensure the command is at the front of the latches.
    pub fn release(&self, lock: &Lock, who: u64) -> Vec<u64> {
        let mut wakeup_list: Vec<u64> = vec![];
        for i in &lock.required_slots[..lock.owned_count] {
            let mut latch = self.slots[*i].lock().unwrap();
            let front = latch.waiting.pop_front().unwrap();
            assert_eq!(front, who);
            if let Some(wakeup) = latch.waiting.front() {
                wakeup_list.push(*wakeup);
            }
            // For some hot keys, the waiting list maybe very long, so we should shrink the waiting
            // VecDeque after pop.
            if latch.waiting.capacity() > WAITING_LIST_SHRINK_SIZE
                && latch.waiting.len() < WAITING_LIST_SHRINK_SIZE
            {
                latch.waiting.shrink_to_fit();
            }
        }
        wakeup_list
    }

    /// Calculates the slot ID by hashing the `key`.
    fn calc_slot<H>(&self, key: &H) -> usize
    where
        H: Hash,
    {
        let mut s = DefaultHasher::new();
        key.hash(&mut s);
        (s.finish() as usize) & (self.size - 1)
    }
}

#[cfg(test)]
mod tests {
    use super::{Latches, Lock};

    #[test]
    fn test_wakeup() {
        let latches = Latches::new(256);

        let slots_a: Vec<usize> = vec![1, 3, 5];
        let mut lock_a = Lock::new(slots_a);
        let slots_b: Vec<usize> = vec![4, 5, 6];
        let mut lock_b = Lock::new(slots_b);
        let cid_a: u64 = 1;
        let cid_b: u64 = 2;

        // a acquire lock success
        let acquired_a = latches.acquire(&mut lock_a, cid_a);
        assert_eq!(acquired_a, true);

        // b acquire lock failed
        let mut acquired_b = latches.acquire(&mut lock_b, cid_b);
        assert_eq!(acquired_b, false);

        // a release lock, and get wakeup list
        let wakeup = latches.release(&lock_a, cid_a);
        assert_eq!(wakeup[0], cid_b);

        // b acquire lock success
        acquired_b = latches.acquire(&mut lock_b, cid_b);
        assert_eq!(acquired_b, true);
    }

    #[test]
    fn test_wakeup_by_multi_cmds() {
        let latches = Latches::new(256);

        let slots_a: Vec<usize> = vec![1, 2, 3];
        let slots_b: Vec<usize> = vec![4, 5, 6];
        let slots_c: Vec<usize> = vec![3, 4];
        let mut lock_a = Lock::new(slots_a);
        let mut lock_b = Lock::new(slots_b);
        let mut lock_c = Lock::new(slots_c);
        let cid_a: u64 = 1;
        let cid_b: u64 = 2;
        let cid_c: u64 = 3;

        // a acquire lock success
        let acquired_a = latches.acquire(&mut lock_a, cid_a);
        assert_eq!(acquired_a, true);

        // b acquire lock success
        let acquired_b = latches.acquire(&mut lock_b, cid_b);
        assert_eq!(acquired_b, true);

        // c acquire lock failed, cause a occupied slot 3
        let mut acquired_c = latches.acquire(&mut lock_c, cid_c);
        assert_eq!(acquired_c, false);

        // a release lock, and get wakeup list
        let wakeup = latches.release(&lock_a, cid_a);
        assert_eq!(wakeup[0], cid_c);

        // c acquire lock failed again, cause b occupied slot 4
        acquired_c = latches.acquire(&mut lock_c, cid_c);
        assert_eq!(acquired_c, false);

        // b release lock, and get wakeup list
        let wakeup = latches.release(&lock_b, cid_b);
        assert_eq!(wakeup[0], cid_c);

        // finally c acquire lock success
        acquired_c = latches.acquire(&mut lock_c, cid_c);
        assert_eq!(acquired_c, true);
    }
}
<|endoftext|>