<commit_before>#![deny(missing_docs)]

//! API that wraps the pandoc command line tool

#[macro_use]
extern crate itertools;

use itertools::Itertools;

/// path to pandoc executable
#[cfg(windows)]
const PANDOC_PATH: &'static [&'static str] = &[
    // this compiles the user's name into the binary, maybe not the greatest idea?
    concat!(env!("LOCALAPPDATA"), r#"\Pandoc\"#),
];
/// path to pandoc executable
#[cfg(not(windows))]
const PANDOC_PATH: &'static [&'static str] = &[
];

/// path where miktex executables can be found
#[cfg(windows)]
const LATEX_PATH: &'static [&'static str] = &[
    r#"C:\Program Files (x86)\MiKTeX 2.9\miktex\bin"#,
    r#"C:\Program Files\MiKTeX 2.9\miktex\bin"#,
];
/// path where miktex executables can be found
#[cfg(not(windows))]
const LATEX_PATH: &'static [&'static str] = &[
    r"/usr/local/bin",
    r"/usr/local/texlive/2015/bin/i386-linux",
];

use std::process::Command;
use std::env;

/// equivalent to the latex document class
#[derive(Debug)]
pub enum DocumentClass {
    /// compact form of report
    Article,
    /// abstract + chapters + custom page for title, abstract and toc
    Report,
    /// no abstract
    Book,
}

pub use DocumentClass::*;

impl DocumentClass {
    fn pandoc_name(&self) -> &'static str {
        match *self {
            Article => "article",
            Report => "report",
            Book => "book",
        }
    }
}

/// the argument builder
#[derive(Debug, Default)]
pub struct Pandoc {
    inputs: Vec<String>,
    output: Option<String>,
    latex_path_hint: Vec<String>,
    pandoc_path_hint: Vec<String>,
    document_class: Option<DocumentClass>,
    bibliography: Option<String>,
    csl: Option<String>,
    toc: bool,
    number_sections: bool,
    template: Option<String>,
    variables: Vec<(String, String)>,
}

use std::convert::Into;
use std::borrow::Cow;

/// does nothing useful, simply gives you a builder object
/// convenience function so you can call pandoc::new()
pub fn new() -> Pandoc { Default::default() }

impl Pandoc {
    /// does nothing useful, simply gives you a builder object
    pub fn new() -> Pandoc { Default::default() }
    /// this path is searched first for latex, then PATH, then some hardcoded hints
    pub fn add_latex_path_hint<'a, T: Into<Cow<'a, str>>>(&mut self, path: T) {
        self.latex_path_hint.push(path.into().into_owned());
    }
    /// this path is searched first for pandoc, then PATH, then some hardcoded hints
    pub fn add_pandoc_path_hint<'a, T: Into<Cow<'a, str>>>(&mut self, path: T) {
        self.pandoc_path_hint.push(path.into().into_owned());
    }

    /// sets or overwrites the document-class
    pub fn set_doc_class(&mut self, class: DocumentClass) {
        self.document_class = Some(class);
    }

    /// adds more input files, the order is relevant
    /// the order of adding the files is the order in which they are processed
    pub fn add_input<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.inputs.push(filename.into().into_owned());
    }
    /// sets or overwrites the output filename
    pub fn set_output<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.output = Some(filename.into().into_owned());
    }

    /// filename of the bibliography database
    pub fn set_bibliography<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.bibliography = Some(filename.into().into_owned());
    }
    /// filename of a citation style file
    pub fn set_csl<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.csl = Some(filename.into().into_owned());
    }
    /// enable table of contents
    pub fn set_toc(&mut self) { self.toc = true; }
    /// prefix section names with indices x.y.z
    pub fn set_number_sections(&mut self) { self.number_sections = true; }
    /// set a custom latex template
    pub fn set_latex_template<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.template = Some(filename.into().into_owned());
    }
    /// set a custom variable
    /// try not to use this, there are convenience functions for most things
    pub fn set_variable<'a, T: Into<Cow<'a, str>>, U: Into<Cow<'a, str>>>(&mut self, key: T, value: U) {
        self.variables.push((key.into().into_owned(), value.into().into_owned()));
    }

    /// actually execute pandoc
    pub fn execute(self) -> Result<(), PandocError> {
        let mut cmd = Command::new("pandoc");
        for input in self.inputs {
            cmd.arg(input);
        }
        cmd.arg("-o").arg(self.output.expect("no output file specified"));
        if let Some(filename) = self.bibliography {
            cmd.arg(format!("--bibliography={}", filename));
        }
        if let Some(filename) = self.csl {
            cmd.arg(format!("--csl={}", filename));
        }
        if self.toc {
            cmd.arg("--toc");
        }
        if self.number_sections {
            cmd.arg("--number-sections");
        }
        if let Some(filename) = self.template {
            cmd.arg(format!("--template={}", filename));
        }
        for (key, val) in self.variables {
            cmd.arg("--variable").arg(format!("{}={}", key, val));
        }
        if let Some(doc_class) = self.document_class {
            cmd.arg("--variable").arg(format!("documentclass={}", doc_class.pandoc_name()));
        }
        let path: String = self.latex_path_hint.iter()
            .chain(self.pandoc_path_hint.iter())
            .map(std::borrow::Borrow::borrow)
            .chain(PANDOC_PATH.into_iter().cloned())
            .chain(LATEX_PATH.into_iter().cloned())
            .chain([env::var("PATH").unwrap()].iter().map(std::borrow::Borrow::borrow))
            .intersperse(";")
            .collect();
        cmd.env("PATH", path);

        match cmd.output() {
            Err(e) => Err(PandocError::IoErr(e)),
            Ok(ref o) if o.status.success() => Ok(()),
            Ok(o) => Err(PandocError::Err(o)),
        }
    }
}

/// Possible errors that can occur before or during pandoc execution
pub enum PandocError {
    /// some kind of IO-Error
    IoErr(std::io::Error),
    /// pandoc execution failed, look at the stderr output
    Err(std::process::Output),
}

impl std::fmt::Debug for PandocError {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        match *self {
            PandocError::IoErr(ref e) => write!(fmt, "{:?}", e),
            PandocError::Err(ref e) => {
                try!(write!(fmt, "exit_code: {:?}", e.status.code()));
                try!(write!(fmt, "stdout: {}", String::from_utf8_lossy(&e.stdout)));
                write!(fmt, "stderr: {}", String::from_utf8_lossy(&e.stderr))
            }
        }
    }
}
<commit_msg>add more options and allow the generation of a latex template from the settings<commit_after>#![deny(missing_docs)]

//! API that wraps the pandoc command line tool

#[macro_use]
extern crate itertools;

use itertools::Itertools;

use std::io::Write;

/// path to pandoc executable
#[cfg(windows)]
const PANDOC_PATH: &'static [&'static str] = &[
    // this compiles the user's name into the binary, maybe not the greatest idea?
    concat!(env!("LOCALAPPDATA"), r#"\Pandoc\"#),
];
/// path to pandoc executable
#[cfg(not(windows))]
const PANDOC_PATH: &'static [&'static str] = &[
];

/// path where miktex executables can be found
#[cfg(windows)]
const LATEX_PATH: &'static [&'static str] = &[
    r#"C:\Program Files (x86)\MiKTeX 2.9\miktex\bin"#,
    r#"C:\Program Files\MiKTeX 2.9\miktex\bin"#,
];
/// path where miktex executables can be found
#[cfg(not(windows))]
const LATEX_PATH: &'static [&'static str] = &[
    r"/usr/local/bin",
    r"/usr/local/texlive/2015/bin/i386-linux",
];

use std::process::Command;
use std::env;

/// equivalent to the latex document class
#[derive(Debug)]
pub enum DocumentClass {
    /// compact form of report
    Article,
    /// abstract + chapters + custom page for title, abstract and toc
    Report,
    /// no abstract
    Book,
}

pub use DocumentClass::*;

impl DocumentClass {
    fn pandoc_name(&self) -> &'static str {
        match *self {
            Article => "article",
            Report => "report",
            Book => "book",
        }
    }
}

/// typesafe access to -t FORMAT, -w FORMAT, --to=FORMAT, --write=FORMAT
#[derive(Debug)]
pub enum OutputFormat {
    /// native Haskell
    Native,
    /// JSON version of native AST
    Json,
    /// Plain text
    Plain,
    /// pandoc’s extended markdown
    Markdown,
    /// original unextended markdown
    MarkdownStrict,
    /// PHP Markdown extra extended markdown
    MarkdownPhpextra,
    /// github extended markdown
    MarkdownGithub,
    /// CommonMark markdown
    Commonmark,
    /// reStructuredText
    Rst,
    /// XHTML 1
    Html,
    /// HTML 5
    Html5,
    /// LaTeX
    Latex,
    /// LaTeX beamer slide show
    Beamer,
    /// ConTeXt
    Context,
    /// Groff man
    Man,
    /// MediaWiki markup
    MediaWiki,
    /// DokuWiki markup
    Dokuwiki,
    /// Textile
    Textile,
    /// Emacs Org-Mode
    Org,
    /// GNU Texinfo
    Texinfo,
    /// OPML
    Opml,
    /// DocBook
    Docbook,
    /// Open Document
    OpenDocument,
    /// OpenOffice text document
    Odt,
    /// Word docx
    Docx,
    /// Haddock markup
    Haddock,
    /// Rich text format
    Rtf,
    /// EPUB v2 book
    Epub,
    /// EPUB v3
    Epub3,
    /// FictionBook2 e-book
    Fb2,
    /// AsciiDoc
    Asciidoc,
    /// InDesign ICML
    Icml,
    /// Slidy HTML and javascript slide show
    Slidy,
    /// Slideous HTML and javascript slide show
    Slideous,
    /// DZSlides HTML5 + javascript slide show
    Dzslides,
    /// reveal.js HTML5 + javascript slide show
    Revealjs,
    /// S5 HTML and javascript slide show
    S5,
    /// the path of a custom lua writer (see Custom writers)
    Lua(String),
}

impl OutputFormat {
    fn pandoc_name(&self) -> &'static str {
        use OutputFormat::*;
        match *self {
            Native => "native",
            Json => "json",
            Plain => "plain",
            Markdown => "markdown",
            MarkdownStrict => "markdown_strict",
            MarkdownPhpextra => "markdown_phpextra",
            MarkdownGithub => "markdown_github",
            Commonmark => "commonmark",
            Rst => "rst",
            Html => "html",
            Html5 => "html5",
            Latex => "latex",
            Beamer => "beamer",
            Context => "context",
            Man => "man",
            MediaWiki => "mediawiki",
            Dokuwiki => "dokuwiki",
            Textile => "textile",
            Org => "org",
            Texinfo => "texinfo",
            Opml => "opml",
            Docbook => "docbook",
            OpenDocument => "open_document",
            Odt => "odt",
            Docx => "docx",
            Haddock => "haddock",
            Rtf => "rtf",
            Epub => "epub",
            Epub3 => "epub3",
            Fb2 => "fb2",
            Asciidoc => "asciidoc",
            Icml => "icml",
            Slidy => "slidy",
            Slideous => "slideous",
            Dzslides => "dzslides",
            Revealjs => "revealjs",
            S5 => "s5",
            Lua(_) => unimplemented!(),
        }
    }
}

/// typesafe access to -f FORMAT, -r FORMAT, --from=FORMAT, --read=FORMAT
#[derive(Debug)]
pub enum InputFormat {
    /// native Haskell
    Native,
    /// JSON version of native AST
    Json,
    /// pandoc’s extended markdown
    Markdown,
    /// original unextended markdown
    MarkdownStrict,
    /// PHP Markdown extra extended markdown
    MarkdownPhpextra,
    /// github extended markdown
    MarkdownGithub,
    /// CommonMark markdown
    Commonmark,
    /// Textile
    Textile,
    /// reStructuredText
    Rst,
    /// HTML
    Html,
    /// DocBook
    DocBook,
    /// txt2tags
    T2t,
    /// Word docx
    Docx,
    /// EPUB
    Epub,
    /// OPML
    Opml,
    /// Emacs Org-Mode
    Org,
    /// MediaWiki markup
    MediaWiki,
    /// TWiki markup
    Twiki,
    /// Haddock markup
    Haddock,
    /// LaTeX
    Latex,
}

#[allow(missing_docs)]
pub enum MarkdownExtensions {
    EscapedLineBreaks,
    BlankBeforeHeader,
    HeaderAttributes,
    AutoIdentifiers,
    ImplicitHeaderReferences,
    BlankBeforeBlockQuote,
    FencedCodeBlocks,
    BacktickCodeBlocks,
    FencedCodeAttributes,
    LineBlocks,
    FancyLists,
    Startnum,
    DefinitionLists,
    ExampleLists,
    TableCaptions,
    SimpleTables,
    MultilineTables,
    GridTables,
    PipeTables,
    PandocTitleBlock,
    YamlMetadataBlock,
    AllSymbolsEscapable,
    IntrawordUnderscores,
    Strikeout,
    Superscript,
    Subscript,
    InlineCodeAttributes,
    TexMathDollars,
    RawHtml,
    MarkdownInHtmlBlocks,
    NativeDivs,
    NativeSpans,
    RawTex,
    LatexMacros,
    ShortcutReferenceLinks,
    ImplicitFigures,
    Footnotes,
    InlineNotes,
    Citations,
    ListsWithoutPrecedingBlankline,
    HardLineBreaks,
    IgnoreLineBreaks,
    TexMathSingleBackslash,
    TexMathDoubleBackslash,
    MarkdownAttribute,
    MmdTitleBlock,
    Abbreviations,
    AutolinkBareUris,
    AsciiIdentifiers,
    LinkAttributes,
    MmdHeaderIdentifiers,
    CompactDefinitionLists,
}

/// the argument builder
#[derive(Debug, Default)]
pub struct Pandoc {
    inputs: Vec<String>,
    output: Option<String>,
    output_format: Option<OutputFormat>,
    latex_path_hint: Vec<String>,
    pandoc_path_hint: Vec<String>,
    document_class: Option<DocumentClass>,
    bibliography: Option<String>,
    csl: Option<String>,
    toc: bool,
    chapters: bool,
    number_sections: bool,
    template: Option<String>,
    variables: Vec<(String, String)>,
    slide_level: Option<usize>,
}

use std::convert::Into;
use std::borrow::Cow;

/// does nothing useful, simply gives you a builder object
/// convenience function so you can call pandoc::new()
pub fn new() -> Pandoc { Default::default() }

impl Pandoc {
    /// does nothing useful, simply gives you a builder object
    pub fn new() -> Pandoc { Default::default() }
    /// this path is searched first for latex, then PATH, then some hardcoded hints
    pub fn add_latex_path_hint<'a, T: Into<Cow<'a, str>>>(&mut self, path: T) {
        self.latex_path_hint.push(path.into().into_owned());
    }
    /// this path is searched first for pandoc, then PATH, then some hardcoded hints
    pub fn add_pandoc_path_hint<'a, T: Into<Cow<'a, str>>>(&mut self, path: T) {
        self.pandoc_path_hint.push(path.into().into_owned());
    }

    /// sets or overwrites the document-class
    pub fn set_doc_class(&mut self, class: DocumentClass) {
        self.document_class = Some(class);
    }
    /// sets or overwrites the output format
    pub fn set_output_format(&mut self, format: OutputFormat) {
        self.output_format = Some(format);
    }

    /// adds more input files, the order is relevant
    /// the order of adding the files is the order in which they are processed
    pub fn add_input<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.inputs.push(filename.into().into_owned());
    }
    /// sets or overwrites the output filename
    pub fn set_output<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.output = Some(filename.into().into_owned());
    }

    /// filename of the bibliography database
    pub fn set_bibliography<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.bibliography = Some(filename.into().into_owned());
    }
    /// filename of a citation style file
    pub fn set_csl<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.csl = Some(filename.into().into_owned());
    }
    /// enable table of contents
    pub fn set_toc(&mut self) { self.toc = true; }
    /// enable chapters
    pub fn set_chapters(&mut self) { self.chapters = true; }
    /// prefix section names with indices x.y.z
    pub fn set_number_sections(&mut self) { self.number_sections = true; }
    /// set a custom latex template
    pub fn set_latex_template<'a, T: Into<Cow<'a, str>>>(&mut self, filename: T) {
        self.template = Some(filename.into().into_owned());
    }
    /// sets the header level that causes a new slide to be generated
    pub fn set_slide_level(&mut self, level: usize) {
        self.slide_level = Some(level);
    }
    /// set a custom variable
    /// try not to use this, there are convenience functions for most things
    pub fn set_variable<'a, T: Into<Cow<'a, str>>, U: Into<Cow<'a, str>>>(&mut self, key: T, value: U) {
        self.variables.push((key.into().into_owned(), value.into().into_owned()));
    }

    fn cmd(self) -> Command {
        let mut cmd = Command::new("pandoc");
        for input in self.inputs {
            cmd.arg(input);
        }
        cmd.arg("-o").arg(self.output.expect("no output file specified"));
        if let Some(format) = self.output_format {
            cmd.arg(format!("--write={}", format.pandoc_name()));
        }
        if let Some(filename) = self.bibliography {
            cmd.arg(format!("--bibliography={}", filename));
        }
        if let Some(filename) = self.csl {
            cmd.arg(format!("--csl={}", filename));
        }
        if self.toc {
            cmd.arg("--toc");
        }
        if self.number_sections {
            cmd.arg("--number-sections");
        }
        if let Some(filename) = self.template {
            cmd.arg(format!("--template={}", filename));
        }
        for (key, val) in self.variables {
            cmd.arg("--variable").arg(format!("{}={}", key, val));
        }
        if let Some(doc_class) = self.document_class {
            cmd.arg("--variable").arg(format!("documentclass={}", doc_class.pandoc_name()));
        }
        if let Some(level) = self.slide_level {
            cmd.arg(format!("--slide-level={}", level));
        }
        let path: String = self.latex_path_hint.iter()
            .chain(self.pandoc_path_hint.iter())
            .map(std::borrow::Borrow::borrow)
            .chain(PANDOC_PATH.into_iter().cloned())
            .chain(LATEX_PATH.into_iter().cloned())
            .chain([env::var("PATH").unwrap()].iter().map(std::borrow::Borrow::borrow))
            .intersperse(";")
            .collect();
        cmd.env("PATH", path);
        cmd
    }

    /// generate a latex template from the given settings
    /// this function can panic in a lot of places
    pub fn generate_latex_template<'a, T: Into<Cow<'a, str>>>(self, filename: T) {
        let format = self.output_format.as_ref().map(OutputFormat::pandoc_name).unwrap();
        let mut cmd = self.cmd();
        cmd.arg(format!("--print-default-template={}", format));
        let output = cmd.output().unwrap();
        assert_eq!(output.status.code().unwrap(), 0);
        let filename: &str = &filename.into();
        let mut file = std::fs::File::create(filename).unwrap();
        file.write(&output.stdout).unwrap();
    }

    /// actually execute pandoc
    pub fn execute(self) -> Result<(), PandocError> {
        let mut cmd = self.cmd();
        match cmd.output() {
            Err(e) => Err(PandocError::IoErr(e)),
            Ok(ref o) if o.status.success() => Ok(()),
            Ok(o) => Err(PandocError::Err(o)),
        }
    }
}

/// Possible errors that can occur before or during pandoc execution
pub enum PandocError {
    /// some kind of IO-Error
    IoErr(std::io::Error),
    /// pandoc execution failed, look at the stderr output
    Err(std::process::Output),
}

impl std::fmt::Debug for PandocError {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        match *self {
            PandocError::IoErr(ref e) => write!(fmt, "{:?}", e),
            PandocError::Err(ref e) => {
                try!(write!(fmt, "exit_code: {:?}", e.status.code()));
                try!(write!(fmt, "stdout: {}", String::from_utf8_lossy(&e.stdout)));
                write!(fmt, "stderr: {}", String::from_utf8_lossy(&e.stderr))
            }
        }
    }
}
<|endoftext|>


<commit_before>//!
//!
//! A library for creating interactive command line shells
#[macro_use] extern crate prettytable;
use prettytable::Table;
use prettytable::format;

use std::io;
use std::io::prelude::*;
use std::string::ToString;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex};

use std::collections::BTreeMap;

/// Command execution error
#[derive(Debug)]
pub enum ExecError {
    /// Empty command provided
    Empty,
    /// Exit from the shell loop
    Quit,
    /// Some arguments are missing
    MissingArgs,
    /// The provided command is unknown
    UnknownCommand(String),
    /// The history index is not valid
    InvalidHistory(usize),
    /// Other error that may have happen during command execution
    Other(Box<Error>),
}
use ExecError::*;

impl fmt::Display for ExecError {
    fn fmt(&self, format: &mut fmt::Formatter) -> fmt::Result {
        return match self {
            &Empty => write!(format, "No command provided"),
            &Quit => write!(format, "Quit"),
            &UnknownCommand(ref cmd) => write!(format, "Unknown Command {}", cmd),
            &InvalidHistory(i) => write!(format, "Invalid history entry {}", i),
            &MissingArgs => write!(format, "Not enough arguments"),
            &Other(ref e) => write!(format, "{}", e)
        };
    }
}

// impl Error for ExecError {
//     fn description(&self) -> &str {
//         return match self {
//             &Quit => "The command requested to quit",
//             &UnknownCommand(..) => "The provided command is unknown",
//             &MissingArgs => "Not enough arguments have been provided",
//             &Other(..) => "Other error occured"
//         };
//     }
// }

impl <E: Error + 'static> From<E> for ExecError {
    fn from(e: E) -> ExecError {
        return Other(Box::new(e));
    }
}

/// Input / Output for shell execution
#[derive(Clone)]
pub struct ShellIO {
    input: Arc<Mutex<io::Read + Send>>,
    output: Arc<Mutex<io::Write + Send>>
}

impl ShellIO {
    /// Create a new Shell I/O wrapping provided Input and Output
    pub fn new(input: I, output: O) -> ShellIO
        where I: Read + Send + 'static, O: Write + Send + 'static
    {
        return ShellIO {
            input: Arc::new(Mutex::new(input)),
            output: Arc::new(Mutex::new(output))
        };
    }

    /// Create a new Shell I/O wrapping provided Read/Write io
    pub fn new_io<T>(io: T) -> ShellIO
        where T: Read + Write + Send + 'static
    {
        let io = Arc::new(Mutex::new(io));
        return ShellIO {
            input: io.clone(),
            output: io
        };
    }
}

impl Read for ShellIO {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        return self.input.lock().expect("Cannot get handle to console input").read(buf);
    }
}

impl Write for ShellIO {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        return self.output.lock().expect("Cannot get handle to console output").write(buf);
    }

    fn flush(&mut self) -> io::Result<()> {
        return self.output.lock().expect("Cannot get handle to console output").flush();
    }
}

impl Default for ShellIO {
    fn default() -> Self {
        return Self::new(io::stdin(), io::stdout());
    }
}

/// Result from command execution
pub type ExecResult = Result<(), ExecError>;

/// A shell
pub struct Shell<T> {
    commands: BTreeMap<String, Arc<builtins::Command<T>>>,
    io: ShellIO,
    data: T,
    prompt: String,
    history: History
}

impl <T> Shell<T> {
    /// Create a new shell, wrapping `data`, using provided IO
    pub fn new_io(data: T, io: ShellIO) -> Shell<T> {
        let mut sh = Shell {
            commands: BTreeMap::new(),
            io: io,
            data: data,
            prompt: String::from(">"),
            history: History::new(10),
        };
        sh.register_command(builtins::help_cmd());
        sh.register_command(builtins::quit_cmd());
        sh.register_command(builtins::history_cmd());
        return sh;
    }

    /// Create a new shell, wrapping `data`, using standard input/outpu
    pub fn new(data: T) -> Shell<T> {
        return Shell::new_io(data, Default::default());
    }

    /// Get a mutable pointer to the inner data
    pub fn data(&mut self) -> &mut T {
        return &mut self.data;
    }

    /// Change the current prompt
    pub fn set_prompt(&mut self, prompt: String) {
        self.prompt = prompt;
    }

    ///Configure Input / Output for shell
    pub fn set_io(&mut self, io: ShellIO) {
        self.io = io;
    }

    fn register_command(&mut self, cmd: builtins::Command<T>) {
        self.commands.insert(cmd.name.clone(), Arc::new(cmd));
    }

    /// Register a shell command.
    /// Shell commands get called with a reference to the current shell
    pub fn new_shell_command<S, F>(&mut self, name: S, description: S, nargs: usize, func: F)
        where S: ToString, F: Fn(&mut Shell<T>, &[&str]) -> ExecResult + Send + Sync + 'static
    {
        self.register_command(builtins::Command::new(name.to_string(), description.to_string(), nargs, Box::new(func)));
    }

    /// Register a command
    pub fn new_command<S, F>(&mut self, name: S, description: S, nargs: usize, func: F)
        where S: ToString, F: Fn(ShellIO, &mut T, &[&str]) -> ExecResult + Send + Sync + 'static
    {
        self.new_shell_command(name, description, nargs, move |sh, args| func(sh.io.clone(), sh.data(), args));
    }

    /// Register a command that do not accept any argument
    pub fn new_command_noargs<S, F>(&mut self, name: S, description: S, func: F)
        where S: ToString, F: Fn(ShellIO, &mut T) -> ExecResult + Send + Sync + 'static
    {
        self.new_shell_command(name, description, 0, move |sh, _| func(sh.io.clone(), sh.data()));
    }

    /// Print the help to stdout
    pub fn print_help(&mut self) -> ExecResult {
        let mut table = Table::new();
        table.set_format(*format::consts::FORMAT_CLEAN);
        for cmd in self.commands.values() {
            table.add_row(cmd.help());
        }
        return table.print(&mut self.io).map_err(|e| From::from(e))
    }

    /// Return the command history
    pub fn get_history(&mut self) -> &mut History {
        return &mut self.history;
    }

    /// Evaluate a command line
    pub fn eval(&mut self, line: &str) -> ExecResult {
        let mut splt = line.trim().split_whitespace();
        return match splt.next() {
            None => Err(Empty),
            Some(cmd) => match self.commands.get(cmd).map(|i| i.clone()) {
                None => Err(UnknownCommand(cmd.to_string())),
                Some(c) => c.run(self, &splt.collect::<Vec<&str>>())
            }
        };
    }

    fn print_prompt(&mut self) {
        write!(self.io, "{}", self.prompt).unwrap();
        self.io.flush().unwrap();
    }

    /// Enter the shell main loop, exiting only when
    /// the "quit" command is called
    pub fn run_loop(&mut self) {
        self.print_prompt();
        let stdin = io::BufReader::new(self.io.clone());

        for line in stdin.lines().map(|l| l.unwrap()) {
            if let Err(e) = self.eval(&line) {
                match e {
                    Empty => {},
                    Quit => return,
                    e @ _ => writeln!(self.io, "Error : {}", e).unwrap()
                };
            } else {
                self.get_history().push(line);
            }
            self.print_prompt();
        }
    }
}

impl <T> Deref for Shell<T> {
    type Target = T;
    fn deref(&self) -> &T {
        return &self.data;
    }
}

impl <T> DerefMut for Shell<T> {
    fn deref_mut(&mut self) -> &mut T {
        return &mut self.data;
    }
}

impl <T> Clone for Shell<T> where T: Clone {
    fn clone(&self) -> Self {
        return Shell {
            commands: self.commands.clone(),
            io: self.io.clone(),
            data: self.data.clone(),
            prompt: self.prompt.clone(),
            history: self.history.clone()
        };
    }
}

/// Wrap the command histroy from a shell.
/// It has a maximum capacity, and when max capacity is reached,
/// less recent command is removed from history
#[derive(Clone)]
pub struct History {
    history: Arc<Mutex<Vec<String>>>,
    capacity: usize
}

impl History {
    /// Create a new history with the given capacity
    fn new(capacity: usize) -> History {
        return History {
            history: Arc::new(Mutex::new(Vec::with_capacity(capacity))),
            capacity: capacity
        };
    }

    /// Push a command to the history, removing the oldest
    /// one if maximum capacity has been reached
    fn push(&mut self, cmd: String) {
        let mut hist = self.history.lock().unwrap();
        if hist.len() >= self.capacity {
            hist.remove(0);
        }
        hist.push(cmd);
    }

    /// Print the history to stdout
    pub fn print<T: Write>(&self, out: &mut T) {
        let mut cnt = 0;
        for s in &*self.history.lock().unwrap() {
            writeln!(out, "{}: {}", cnt, s).expect("Cannot write to output");
            cnt += 1;
        }
    }

    /// Get a command from history by its index
    pub fn get(&self, i: usize) -> Option<String> {
        return self.history.lock().unwrap().get(i).map(|s| s.clone());
    }
}

mod builtins {
    use std::str::FromStr;
    use prettytable::row::Row;
    use super::{Shell, ExecError, ExecResult};

    pub type CmdFn<T> = Box<Fn(&mut Shell<T>, &[&str]) -> ExecResult + Send + Sync>;

    pub struct Command<T> {
        pub name: String,
        description: String,
        nargs: usize,
        func: CmdFn<T>
    }

    impl <T> Command<T> {
        pub fn new(name: String, description: String, nargs: usize, func: CmdFn<T>) -> Command<T> {
            return Command {
                name: name,
                description: description,
                nargs: nargs,
                func: func
            };
        }

        pub fn help(&self) -> Row {
            return row![self.name, ":", self.description];
        }

        pub fn run(&self, shell: &mut Shell<T>, args: &[&str]) -> ExecResult {
            if args.len() < self.nargs {
                return Err(ExecError::MissingArgs);
            }
            return (self.func)(shell, args);
        }
    }

    pub fn help_cmd<T>() -> Command<T> {
        return Command::new("help".to_string(), "Print this help".to_string(), 0, Box::new(|shell, _| shell.print_help()));
    }

    pub fn quit_cmd<T>() -> Command<T> {
        return Command::new("quit".to_string(), "Quit".to_string(), 0, Box::new(|_, _| Err(ExecError::Quit)));
    }

    pub fn history_cmd<T>() -> Command<T> {
        return Command::new("history".to_string(), "Print commands history or run a command from it".to_string(), 0, Box::new(|shell, args| {
            if args.len() > 0 {
                let i = try!(usize::from_str(args[0]));
                let cmd = try!(shell.get_history().get(i).ok_or(ExecError::InvalidHistory(i)));
                return shell.eval(&cmd);
            } else {
                let mut io = shell.io.clone();
                shell.get_history().print(&mut io);
                return Ok(());
            }
        }));
    }
}
<commit_msg>Fixed doc typo<commit_after>//!
//!
//! A library for creating interactive command line shells
#[macro_use] extern crate prettytable;
use prettytable::Table;
use prettytable::format;

use std::io;
use std::io::prelude::*;
use std::string::ToString;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex};

use std::collections::BTreeMap;

/// Command execution error
#[derive(Debug)]
pub enum ExecError {
    /// Empty command provided
    Empty,
    /// Exit from the shell loop
    Quit,
    /// Some arguments are missing
    MissingArgs,
    /// The provided command is unknown
    UnknownCommand(String),
    /// The history index is not valid
    InvalidHistory(usize),
    /// Other error that may have happen during command execution
    Other(Box<Error>),
}
use ExecError::*;

impl fmt::Display for ExecError {
    fn fmt(&self, format: &mut fmt::Formatter) -> fmt::Result {
        return match self {
            &Empty => write!(format, "No command provided"),
            &Quit => write!(format, "Quit"),
            &UnknownCommand(ref cmd) => write!(format, "Unknown Command {}", cmd),
            &InvalidHistory(i) => write!(format, "Invalid history entry {}", i),
            &MissingArgs => write!(format, "Not enough arguments"),
            &Other(ref e) => write!(format, "{}", e)
        };
    }
}

// impl Error for ExecError {
//     fn description(&self) -> &str {
//         return match self {
//             &Quit => "The command requested to quit",
//             &UnknownCommand(..) => "The provided command is unknown",
//             &MissingArgs => "Not enough arguments have been provided",
//             &Other(..) => "Other error occured"
//         };
//     }
// }

impl <E: Error + 'static> From<E> for ExecError {
    fn from(e: E) -> ExecError {
        return Other(Box::new(e));
    }
}

/// Input / Output for shell execution
#[derive(Clone)]
pub struct ShellIO {
    input: Arc<Mutex<io::Read + Send>>,
    output: Arc<Mutex<io::Write + Send>>
}

impl ShellIO {
    /// Create a new Shell I/O wrapping provided Input and Output
    pub fn new(input: I, output: O) -> ShellIO
        where I: Read + Send + 'static, O: Write + Send + 'static
    {
        return ShellIO {
            input: Arc::new(Mutex::new(input)),
            output: Arc::new(Mutex::new(output))
        };
    }

    /// Create a new Shell I/O wrapping provided Read/Write io
    pub fn new_io<T>(io: T) -> ShellIO
        where T: Read + Write + Send + 'static
    {
        let io = Arc::new(Mutex::new(io));
        return ShellIO {
            input: io.clone(),
            output: io
        };
    }
}

impl Read for ShellIO {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        return self.input.lock().expect("Cannot get handle to console input").read(buf);
    }
}

impl Write for ShellIO {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        return self.output.lock().expect("Cannot get handle to console output").write(buf);
    }

    fn flush(&mut self) -> io::Result<()> {
        return self.output.lock().expect("Cannot get handle to console output").flush();
    }
}

impl Default for ShellIO {
    fn default() -> Self {
        return Self::new(io::stdin(), io::stdout());
    }
}

/// Result from command execution
pub type ExecResult = Result<(), ExecError>;

/// A shell
pub struct Shell<T> {
    commands: BTreeMap<String, Arc<builtins::Command<T>>>,
    io: ShellIO,
    data: T,
    prompt: String,
    history: History
}

impl <T> Shell<T> {
    /// Create a new shell, wrapping `data`, using provided IO
    pub fn new_io(data: T, io: ShellIO) -> Shell<T> {
        let mut sh = Shell {
            commands: BTreeMap::new(),
            io: io,
            data: data,
            prompt: String::from(">"),
            history: History::new(10),
        };
        sh.register_command(builtins::help_cmd());
        sh.register_command(builtins::quit_cmd());
        sh.register_command(builtins::history_cmd());
        return sh;
    }

    /// Create a new shell, wrapping `data`, using standard input/outpu
    pub fn new(data: T) -> Shell<T> {
        return Shell::new_io(data, Default::default());
    }

    /// Get a mutable pointer to the inner data
    pub fn data(&mut self) -> &mut T {
        return &mut self.data;
    }

    /// Change the current prompt
    pub fn set_prompt(&mut self, prompt: String) {
        self.prompt = prompt;
    }

    ///Configure Input / Output for shell
    pub fn set_io(&mut self, io: ShellIO) {
        self.io = io;
    }

    fn register_command(&mut self, cmd: builtins::Command<T>) {
        self.commands.insert(cmd.name.clone(), Arc::new(cmd));
    }

    /// Register a shell command.
    /// Shell commands get called with a reference to the current shell
    pub fn new_shell_command<S, F>(&mut self, name: S, description: S, nargs: usize, func: F)
        where S: ToString, F: Fn(&mut Shell<T>, &[&str]) -> ExecResult + Send + Sync + 'static
    {
        self.register_command(builtins::Command::new(name.to_string(), description.to_string(), nargs, Box::new(func)));
    }

    /// Register a command
    pub fn new_command<S, F>(&mut self, name: S, description: S, nargs: usize, func: F)
        where S: ToString, F: Fn(ShellIO, &mut T, &[&str]) -> ExecResult + Send + Sync + 'static
    {
        self.new_shell_command(name, description, nargs, move |sh, args| func(sh.io.clone(), sh.data(), args));
    }

    /// Register a command that do not accept any argument
    pub fn new_command_noargs<S, F>(&mut self, name: S, description: S, func: F)
        where S: ToString, F: Fn(ShellIO, &mut T) -> ExecResult + Send + Sync + 'static
    {
        self.new_shell_command(name, description, 0, move |sh, _| func(sh.io.clone(), sh.data()));
    }

    /// Print the help to stdout
    pub fn print_help(&mut self) -> ExecResult {
        let mut table = Table::new();
        table.set_format(*format::consts::FORMAT_CLEAN);
        for cmd in self.commands.values() {
            table.add_row(cmd.help());
        }
        return table.print(&mut self.io).map_err(|e| From::from(e))
    }

    /// Return the command history
    pub fn get_history(&mut self) -> &mut History {
        return &mut self.history;
    }

    /// Evaluate a command line
    pub fn eval(&mut self, line: &str) -> ExecResult {
        let mut splt = line.trim().split_whitespace();
        return match splt.next() {
            None => Err(Empty),
            Some(cmd) => match self.commands.get(cmd).map(|i| i.clone()) {
                None => Err(UnknownCommand(cmd.to_string())),
                Some(c) => c.run(self, &splt.collect::<Vec<&str>>())
            }
        };
    }

    fn print_prompt(&mut self) {
        write!(self.io, "{}", self.prompt).unwrap();
        self.io.flush().unwrap();
    }

    /// Enter the shell main loop, exiting only when
    /// the "quit" command is called
    pub fn run_loop(&mut self) {
        self.print_prompt();
        let stdin = io::BufReader::new(self.io.clone());

        for line in stdin.lines().map(|l| l.unwrap()) {
            if let Err(e) = self.eval(&line) {
                match e {
                    Empty => {},
                    Quit => return,
                    e @ _ => writeln!(self.io, "Error : {}", e).unwrap()
                };
            } else {
                self.get_history().push(line);
            }
            self.print_prompt();
        }
    }
}

impl <T> Deref for Shell<T> {
    type Target = T;
    fn deref(&self) -> &T {
        return &self.data;
    }
}

impl <T> DerefMut for Shell<T> {
    fn deref_mut(&mut self) -> &mut T {
        return &mut self.data;
    }
}

impl <T> Clone for Shell<T> where T: Clone {
    fn clone(&self) -> Self {
        return Shell {
            commands: self.commands.clone(),
            io: self.io.clone(),
            data: self.data.clone(),
            prompt: self.prompt.clone(),
            history: self.history.clone()
        };
    }
}

/// Wrap the command history from a shell.
/// It has a maximum capacity, and when max capacity is reached,
/// less recent command is removed from history
#[derive(Clone)]
pub struct History {
    history: Arc<Mutex<Vec<String>>>,
    capacity: usize
}

impl History {
    /// Create a new history with the given capacity
    fn new(capacity: usize) -> History {
        return History {
            history: Arc::new(Mutex::new(Vec::with_capacity(capacity))),
            capacity: capacity
        };
    }

    /// Push a command to the history, removing the oldest
    /// one if maximum capacity has been reached
    fn push(&mut self, cmd: String) {
        let mut hist = self.history.lock().unwrap();
        if hist.len() >= self.capacity {
            hist.remove(0);
        }
        hist.push(cmd);
    }

    /// Print the history to stdout
    pub fn print<T: Write>(&self, out: &mut T) {
        let mut cnt = 0;
        for s in &*self.history.lock().unwrap() {
            writeln!(out, "{}: {}", cnt, s).expect("Cannot write to output");
            cnt += 1;
        }
    }

    /// Get a command from history by its index
    pub fn get(&self, i: usize) -> Option<String> {
        return self.history.lock().unwrap().get(i).map(|s| s.clone());
    }
}

mod builtins {
    use std::str::FromStr;
    use prettytable::row::Row;
    use super::{Shell, ExecError, ExecResult};

    pub type CmdFn<T> = Box<Fn(&mut Shell<T>, &[&str]) -> ExecResult + Send + Sync>;

    pub struct Command<T> {
        pub name: String,
        description: String,
        nargs: usize,
        func: CmdFn<T>
    }

    impl <T> Command<T> {
        pub fn new(name: String, description: String, nargs: usize, func: CmdFn<T>) -> Command<T> {
            return Command {
                name: name,
                description: description,
                nargs: nargs,
                func: func
            };
        }

        pub fn help(&self) -> Row {
            return row![self.name, ":", self.description];
        }

        pub fn run(&self, shell: &mut Shell<T>, args: &[&str]) -> ExecResult {
            if args.len() < self.nargs {
                return Err(ExecError::MissingArgs);
            }
            return (self.func)(shell, args);
        }
    }

    pub fn help_cmd<T>() -> Command<T> {
        return Command::new("help".to_string(), "Print this help".to_string(), 0, Box::new(|shell, _| shell.print_help()));
    }

    pub fn quit_cmd<T>() -> Command<T> {
        return Command::new("quit".to_string(), "Quit".to_string(), 0, Box::new(|_, _| Err(ExecError::Quit)));
    }

    pub fn history_cmd<T>() -> Command<T> {
        return Command::new("history".to_string(), "Print commands history or run a command from it".to_string(), 0, Box::new(|shell, args| {
            if args.len() > 0 {
                let i = try!(usize::from_str(args[0]));
                let cmd = try!(shell.get_history().get(i).ok_or(ExecError::InvalidHistory(i)));
                return shell.eval(&cmd);
            } else {
                let mut io = shell.io.clone();
                shell.get_history().print(&mut io);
                return Ok(());
            }
        }));
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

struct LSystem {
    alphabet: Vec<char>,
    constants: Vec<char>,
    start: String,
    rules: HashMap<char, String>,
}

impl LSystem {
    /// Some guarantees of a LSystem:
    /// 1. Every char in start is in alphabet
    /// 2. Given a rule X -> Y, X and Y are in alphabet
    /// 3. Each char in alphabet has a rule that maps it to something else
    ///   -> by default a char maps to itself, but this can be overriden with
    ///   the `push` method

    pub fn new(start: String) -> LSystem {
        // push chars in start to alphabet
        // and make them map to themselves in a rule
        let mut alphabet: Vec<char> = Vec::new();
        let mut rules: HashMap<char, String> = HashMap::new();
        for c in start.chars() {
            if !alphabet.contains(&c) {
                alphabet.push(c);
                rules.insert(c, c.to_string());
            }
        }

        LSystem {
            alphabet: alphabet,
            constants: Vec::new(),
            start: start,
            rules: rules,
        }
    }

    pub fn to_string(&self) -> String {
        self.start.clone()
    }

    pub fn next(&mut self) -> LSystem {
        let mut next = String::from("");
        for c in self.start.chars() {
            next.push_str(self.get(c));
        }

        LSystem {
            alphabet: self.alphabet.clone(),
            constants: self.constants.clone(),
            start: next,
            rules: self.rules.clone(),
        }
    }

    fn add_char(&mut self, c: char) {
        if !self.alphabet.contains(&c) {
            self.alphabet.push(c);
            self.rules.insert(c, c.to_string());
        }
    }

    pub fn push(&mut self, rule: HashMap<char, String>) {
        for (key, value) in &rule {
            // Make sure that key, value is in alphabet
            for c in value.chars() {
                self.add_char(c);
            }
            if !self.rules.contains_key(key) {
                self.add_char(*key);
            }
            // Add key value pair
            *self.rules.get_mut(key).unwrap() = value.to_owned();
        }
    }

    /// Returns the value that key maps to
    pub fn get(&self, key: char) -> &String {
        &self.rules[&key]
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use LSystem;

    #[test]
    fn test_new_should_have_correct_alphabet() {
        let lsys: LSystem = LSystem::new("AB".to_string());
        assert!(lsys.alphabet.contains(&'A'));
    }

    #[test]
    fn test_new_should_have_correct_start() {
        let lsys: LSystem = LSystem::new("AB".to_string());
        assert_eq!(lsys.start, "AB");
    }

    #[test]
    fn test_new_should_not_have_duplicate_alphabet_chars() {
        let lsys: LSystem = LSystem::new("AA".to_string());
        assert_eq!(lsys.alphabet.len(), 1);
    }

    #[test]
    fn test_new_should_add_default_rules() {
        let lsys: LSystem = LSystem::new("A".to_string());
        assert_eq!(*lsys.get('A'), "A");
    }

    #[test]
    fn test_push_should_override_existing_rule() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        assert_eq!(*lsys.get('A'), "B");
    }

    #[test]
    fn test_push_should_add_to_alphabet() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        assert!(lsys.alphabet.contains(&'B'));
    }

    #[test]
    fn test_push_should_add_default_rules_for_new_chars() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        assert_eq!(*lsys.get('B'), "B");
    }

    #[test]
    fn test_push_should_work_with_multiple_rules() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "AB".to_string());
        rule.insert('B', "A".to_string());
        lsys.push(rule);
        assert_eq!(*lsys.get('B'), "A");
        assert_eq!(*lsys.get('A'), "AB");
    }

    #[test]
    fn test_next_with_simple_rules_should_have_correct_start() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        let next: LSystem = lsys.next();
        assert_eq!(next.start, "B");
    }

    #[test]
    fn test_next_with_complex_rules_should_have_correct_start() {
        let mut lsys: LSystem = LSystem::new("AB".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "AB".to_string());
        rule.insert('B', "A".to_string());
        lsys.push(rule);
        let next: LSystem = lsys.next();
        assert_eq!(next.start, "ABA");
    }

    #[test]
    fn test_next_should_work_in_chains() {
        let mut lsys: LSystem = LSystem::new("AB".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "AB".to_string());
        rule.insert('B', "A".to_string());
        lsys.push(rule);
        let next: LSystem = lsys.next().next();
        assert_eq!(next.start, "ABAAB");
    }

    #[test]
    fn test_to_string_should_return_start_string() {
        let lsys: LSystem = LSystem::new("AB".to_string());
        assert_eq!(lsys.to_string(), "AB");
    }
}
<commit_msg>Implement constants using add_constant method<commit_after>use std::collections::HashMap;

struct LSystem {
    alphabet: Vec<char>,
    constants: Vec<char>,
    start: String,
    rules: HashMap<char, String>,
}

impl LSystem {
    /// Some guarantees of a LSystem:
    /// 1. Every char in start is in alphabet
    /// 2. Given a rule X -> Y, X and Y are in alphabet
    /// 3. Each char in alphabet has a rule that maps it to something else
    ///   -> by default a char maps to itself, but this can be overwriten with
    ///   the `push` method
    /// See https://en.wikipedia.org/wiki/L-system for more details

    pub fn new(start: String) -> LSystem {
        // push chars in start to alphabet
        // and make them map to themselves in a rule
        let mut alphabet: Vec<char> = Vec::new();
        let mut rules: HashMap<char, String> = HashMap::new();
        for c in start.chars() {
            if !alphabet.contains(&c) {
                alphabet.push(c);
                rules.insert(c, c.to_string());
            }
        }

        LSystem {
            alphabet: alphabet,
            constants: Vec::new(),
            start: start,
            rules: rules,
        }
    }

    pub fn to_string(&self) -> String {
        self.start.clone()
    }

    /// Iterate to the next LSystem string, following all rules
    pub fn next(&mut self) -> LSystem {
        let mut next = String::from("");
        for c in self.start.chars() {
            next.push_str(self.get(c));
        }

        LSystem {
            alphabet: self.alphabet.clone(),
            constants: self.constants.clone(),
            start: next,
            rules: self.rules.clone(),
        }
    }

    /// Add a character to the alphabet with the default 'identity rule'
    pub fn add(&mut self, c: char) {
        if !self.alphabet.contains(&c) {
            self.alphabet.push(c);
            self.rules.insert(c, c.to_string());
        }
    }

    /// Helper function for push
    fn overwrite_rule(&mut self, key: &char, value: &String) {
        *self.rules.get_mut(key).unwrap() = value.to_owned();
    }

    /// Add a constant to the LSystem
    /// Constants always map to themselves and cannot be changed
    pub fn add_constant(&mut self, c: char) {
        self.add(c);
        self.constants.push(c);
    }

    /// Add a new rule to the LSystem
    /// If the values in the rule do not exist, they are added
    pub fn push(&mut self, rule: HashMap<char, String>) {
        for (key, value) in &rule {
            // Make sure key is not a constant
            if !self.constants.contains(&key) {
                for c in value.chars() {
                    self.add(c);
                }
                if !self.rules.contains_key(key) {
                    self.add(*key);
                }
                self.overwrite_rule(key, value);
            }
        }
    }

    /// Returns the rule for which key maps to
    pub fn get(&self, key: char) -> &String {
        &self.rules[&key]
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use LSystem;

    #[test]
    fn test_new_should_have_correct_alphabet() {
        let lsys: LSystem = LSystem::new("AB".to_string());
        assert!(lsys.alphabet.contains(&'A'));
    }

    #[test]
    fn test_new_should_have_correct_start() {
        let lsys: LSystem = LSystem::new("AB".to_string());
        assert_eq!(lsys.start, "AB");
    }

    #[test]
    fn test_new_should_not_have_duplicate_alphabet_chars() {
        let lsys: LSystem = LSystem::new("AA".to_string());
        assert_eq!(lsys.alphabet.len(), 1);
    }

    #[test]
    fn test_new_should_add_default_rules() {
        let lsys: LSystem = LSystem::new("A".to_string());
        assert_eq!(*lsys.get('A'), "A");
    }

    #[test]
    fn test_add_should_add_to_alphabet() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        lsys.add('B');
        assert!(lsys.alphabet.contains(&'B'));
    }

    #[test]
    fn test_add_should_add_default_rule() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        lsys.add('B');
        assert_eq!(*lsys.get('B'), "B");
    }

    #[test]
    fn test_add_constant_should_add_to_constants() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        lsys.add_constant('B');
        assert!(lsys.constants.contains(&'B'));
    }

    #[test]
    fn test_add_constant_should_add_to_alphabet() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        lsys.add_constant('B');
        assert!(lsys.alphabet.contains(&'B'));
    }

    #[test]
    fn test_add_constant_should_add_default_rule() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        lsys.add_constant('B');
        assert_eq!(*lsys.get('B'), "B");
    }

    #[test]
    fn test_push_should_overwrite_existing_rule() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        assert_eq!(*lsys.get('A'), "B");
    }

    #[test]
    fn test_push_should_add_to_alphabet() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        assert!(lsys.alphabet.contains(&'B'));
    }

    #[test]
    fn test_push_should_add_default_rules_for_new_chars() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        assert_eq!(*lsys.get('B'), "B");
    }

    #[test]
    fn test_push_should_not_overwrite_constants() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        lsys.add_constant('A');
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        assert_eq!(*lsys.get('A'), "A");
    }

    #[test]
    fn test_push_should_work_with_multiple_rules() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "AB".to_string());
        rule.insert('B', "A".to_string());
        lsys.push(rule);
        assert_eq!(*lsys.get('B'), "A");
        assert_eq!(*lsys.get('A'), "AB");
    }

    #[test]
    fn test_next_with_simple_rules_should_have_correct_start() {
        let mut lsys: LSystem = LSystem::new("A".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "B".to_string());
        lsys.push(rule);
        let next: LSystem = lsys.next();
        assert_eq!(next.start, "B");
    }

    #[test]
    fn test_next_with_complex_rules_should_have_correct_start() {
        let mut lsys: LSystem = LSystem::new("AB".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "AB".to_string());
        rule.insert('B', "A".to_string());
        lsys.push(rule);
        let next: LSystem = lsys.next();
        assert_eq!(next.start, "ABA");
    }

    #[test]
    fn test_next_should_work_in_chains() {
        let mut lsys: LSystem = LSystem::new("AB".to_string());
        let mut rule: HashMap<char, String> = HashMap::new();
        rule.insert('A', "AB".to_string());
        rule.insert('B', "A".to_string());
        lsys.push(rule);
        let next: LSystem = lsys.next().next();
        assert_eq!(next.start, "ABAAB");
    }

    #[test]
    fn test_to_string_should_return_start_string() {
        let lsys: LSystem = LSystem::new("AB".to_string());
        assert_eq!(lsys.to_string(), "AB");
    }
}
<|endoftext|>


<commit_before>//! A generic connection pool.
//!
//! Opening a new database connection every time one is needed is both
//! inefficient and can lead to resource exhaustion under high traffic
//! conditions. A connection pool maintains a set of open connections to a
//! database, handing them out for repeated use.
//!
//! r2d2 is agnostic to the connection type it is managing. Implementors of the
//! `ManageConnection` trait provide the database-specific logic to create and
//! check the health of connections.
//!
//! # Example
//!
//! Using an imaginary "foodb" database.
//!
//! ```rust,ignore
//! use std::thread;
//!
//! extern crate r2d2;
//! extern crate r2d2_foodb;
//!
//! fn main() {
//!     let config = r2d2::Config::builder()
//!         .pool_size(15)
//!         .build();
//!     let manager = r2d2_foodb::FooConnectionManager::new("localhost:1234");
//!
//!     let pool = r2d2::Pool::new(config, manager).unwrap();
//!
//!     for _ in 0..20 {
//!         let pool = pool.clone();
//!         thread::spawn(move || {
//!             let conn = pool.get().unwrap();
//!             // use the connection
//!             // it will be returned to the pool when it falls out of scope.
//!         })
//!     }
//! }
//! ```
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/r2d2/doc/v0.7.0")]

extern crate antidote;
#[macro_use]
extern crate log;

use antidote::{Mutex, Condvar};
use std::cmp;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::mem;
use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};

#[doc(inline)]
pub use config::Config;

use task::ScheduledThreadPool;

pub mod config;
mod task;
mod thunk;

#[cfg(test)]
mod test;

/// A trait which provides connection-specific functionality.
pub trait ManageConnection: Send + Sync + 'static {
    /// The connection type this manager deals with.
    type Connection: Send + 'static;

    /// The error type returned by `Connection`s.
    type Error: Error + 'static;

    /// Attempts to create a new connection.
    fn connect(&self) -> Result<Self::Connection, Self::Error>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut Self::Connection) -> bool;
}

/// A trait which handles errors reported by the `ManageConnection`.
pub trait HandleError<E>: fmt::Debug + Send + Sync + 'static {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

/// A `HandleError` implementation which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NopErrorHandler;

impl<E> HandleError<E> for NopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// A `HandleError` implementation which logs at the error level.
#[derive(Copy, Clone, Debug)]
pub struct LoggingErrorHandler;

impl<E> HandleError<E> for LoggingErrorHandler
    where E: Error
{
    fn handle_error(&self, error: E) {
        error!("{}", error);
    }
}

/// A trait which allows for customization of connections.
pub trait CustomizeConnection<C, E>: fmt::Debug + Send + Sync + 'static {
    /// Called with connections immediately after they are returned from
    /// `ManageConnection::connect`.
    ///
    /// The default implementation simply returns `Ok(())`.
    ///
    /// # Errors
    ///
    /// If this method returns an error, the connection will be discarded.
    #[allow(unused_variables)]
    fn on_acquire(&self, conn: &mut C) -> Result<(), E> {
        Ok(())
    }
}

/// A `CustomizeConnection` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NopConnectionCustomizer;

impl<C, E> CustomizeConnection<C, E> for NopConnectionCustomizer {}

struct Conn<C> {
    conn: C,
    birth: Instant,
    idle_start: Instant,
}

impl<C> Conn<C> {
    fn new(conn: C) -> Conn<C> {
        let now = Instant::now();
        Conn {
            conn: conn,
            birth: now,
            idle_start: now,
        }
    }
}

struct PoolInternals<C> {
    conns: VecDeque<Conn<C>>,
    num_conns: u32,
    pending_conns: u32,
    last_error: Option<String>,
}

struct SharedPool<M>
    where M: ManageConnection
{
    config: Config<M::Connection, M::Error>,
    manager: M,
    internals: Mutex<PoolInternals<M::Connection>>,
    cond: Condvar,
    thread_pool: ScheduledThreadPool,
}

fn drop_conn<M>(shared: &Arc<SharedPool<M>>, internals: &mut PoolInternals<M::Connection>)
    where M: ManageConnection
{
    internals.num_conns -= 1;

    let min = shared.config.min_idle().unwrap_or(shared.config.pool_size());
    if internals.num_conns + internals.pending_conns < min {
        add_connection(shared, internals);
    }
}

fn add_connection<M>(shared: &Arc<SharedPool<M>>, internals: &mut PoolInternals<M::Connection>)
    where M: ManageConnection
{
    internals.pending_conns += 1;
    inner(Duration::from_secs(0), shared);

    fn inner<M>(delay: Duration, shared: &Arc<SharedPool<M>>)
        where M: ManageConnection
    {
        let new_shared = Arc::downgrade(shared);
        shared.thread_pool.run_after(delay, move || {
            let shared = match new_shared.upgrade() {
                Some(shared) => shared,
                None => return,
            };

            let conn = shared.manager.connect().and_then(|mut conn| {
                shared.config.connection_customizer().on_acquire(&mut conn).map(|_| conn)
            });
            match conn {
                Ok(conn) => {
                    let mut internals = shared.internals.lock();
                    internals.last_error = None;
                    internals.conns.push_back(Conn::new(conn));
                    internals.pending_conns -= 1;
                    internals.num_conns += 1;
                    shared.cond.notify_one();
                }
                Err(err) => {
                    shared.internals.lock().last_error = Some(err.to_string());
                    shared.config.error_handler().handle_error(err);
                    let delay = cmp::max(Duration::from_millis(200), delay);
                    let delay = cmp::min(shared.config.connection_timeout() / 2, delay * 2);
                    inner(delay, &shared);
                }
            }
        });
    }
}

fn reap_connections<M>(shared: &Weak<SharedPool<M>>)
    where M: ManageConnection
{
    let shared = match shared.upgrade() {
        Some(shared) => shared,
        None => return,
    };

    let mut old = VecDeque::with_capacity(shared.config.pool_size() as usize);
    let mut to_drop = vec![];

    let mut internals = shared.internals.lock();
    mem::swap(&mut old, &mut internals.conns);
    let now = Instant::now();
    for conn in old {
        let mut reap = false;
        if let Some(timeout) = shared.config.idle_timeout() {
            reap |= now - conn.idle_start >= timeout;
        }
        if let Some(lifetime) = shared.config.max_lifetime() {
            reap |= now - conn.birth >= lifetime;
        }
        if reap {
            drop_conn(&shared, &mut internals);
            to_drop.push(conn.conn);
        } else {
            internals.conns.push_back(conn);
        }
    }
    drop(internals); // make sure we run to_drop destructors without this locked
}

/// A generic connection pool.
pub struct Pool<M: ManageConnection>(Arc<SharedPool<M>>);

/// Returns a new `Pool` referencing the same state as `self`.
impl<M> Clone for Pool<M>
    where M: ManageConnection
{
    fn clone(&self) -> Pool<M> {
        Pool(self.0.clone())
    }
}

impl<M> fmt::Debug for Pool<M>
    where M: ManageConnection + fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.0.internals.lock();

        fmt.debug_struct("Pool")
           .field("connections", &inner.num_conns)
           .field("idle_connections", &inner.conns.len())
           .field("config", &self.0.config)
           .field("manager", &self.0.manager)
           .finish()
    }
}

impl<M> Pool<M>
    where M: ManageConnection
{
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value if `initialization_fail_fast` is set to true in
    /// the configuration and the pool is unable to open all of its
    /// connections.
    pub fn new(config: Config<M::Connection, M::Error>,
               manager: M)
               -> Result<Pool<M>, InitializationError> {
        Pool::new_inner(config, manager, Duration::from_secs(30))
    }

    // for testing
    fn new_inner(config: Config<M::Connection, M::Error>,
                 manager: M,
                 reaper_rate: Duration)
                 -> Result<Pool<M>, InitializationError> {
        let internals = PoolInternals {
            conns: VecDeque::with_capacity(config.pool_size() as usize),
            num_conns: 0,
            pending_conns: 0,
            last_error: None,
        };

        let shared = Arc::new(SharedPool {
            thread_pool: ScheduledThreadPool::new(config.helper_threads() as usize),
            config: config,
            manager: manager,
            internals: Mutex::new(internals),
            cond: Condvar::new(),
        });

        let initial_size = shared.config.min_idle().unwrap_or(shared.config.pool_size());
        {
            let mut inner = shared.internals.lock();
            for _ in 0..initial_size {
                add_connection(&shared, &mut inner);
            }
            drop(inner);
        }

        if shared.config.initialization_fail_fast() {
            let end = Instant::now() + shared.config.connection_timeout();
            let mut internals = shared.internals.lock();

            while internals.num_conns != initial_size {
                let now = Instant::now();
                if now >= end {
                    return Err(InitializationError(internals.last_error.take()));
                }
                internals = shared.cond
                                  .wait_timeout(internals, end - now)
                                  .0;
            }
        }

        if shared.config.max_lifetime().is_some() || shared.config.idle_timeout().is_some() {
            let s = Arc::downgrade(&shared);
            shared.thread_pool
                  .run_at_fixed_rate(reaper_rate, move || reap_connections(&s));
        }

        Ok(Pool(shared))
    }

    /// Retrieves a connection from the pool.
    ///
    /// Waits for at most `Config::connection_timeout` before returning an
    /// error.
    pub fn get(&self) -> Result<PooledConnection<M>, GetTimeout> {
        let end = Instant::now() + self.0.config.connection_timeout();
        let mut internals = self.0.internals.lock();

        let connection;
        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.0.config.test_on_check_out() {
                        if let Err(e) = self.0.manager.is_valid(&mut conn.conn) {
                            let msg = e.to_string();
                            self.0.config.error_handler().handle_error(e);
                            internals = self.0.internals.lock();
                            internals.last_error = Some(msg);
                            drop_conn(&self.0, &mut internals);
                            continue;
                        }
                    }

                    connection = conn;
                    break;
                }
                None => {
                    if internals.num_conns + internals.pending_conns < self.0.config.pool_size() {
                        add_connection(&self.0, &mut internals);
                    }

                    let now = Instant::now();
                    if now >= end {
                        return Err(GetTimeout(internals.last_error.take()));
                    };
                    internals = self.0
                                    .cond
                                    .wait_timeout(internals, end - now)
                                    .0;
                }
            }
        }

        Ok(PooledConnection {
            pool: self.clone(),
            conn: Some(connection),
        })
    }

    fn put_back(&self, mut conn: Conn<M::Connection>) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.0.manager.has_broken(&mut conn.conn);

        let mut internals = self.0.internals.lock();
        if broken {
            drop_conn(&self.0, &mut internals);
        } else {
            conn.idle_start = Instant::now();
            internals.conns.push_back(conn);
            self.0.cond.notify_one();
        }
    }
}

/// An error returned by `Pool::new` if it fails to initialize connections.
#[derive(Debug)]
pub struct InitializationError(Option<String>);

impl fmt::Display for InitializationError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        try!(fmt.write_str(self.description()));
        if let Some(ref err) = self.0 {
            try!(write!(fmt, ": {}", err));
        }
        Ok(())
    }
}

impl Error for InitializationError {
    fn description(&self) -> &str {
        "Unable to initialize connections"
    }
}

/// An error returned by `Pool::get` if it times out without retrieving a connection.
#[derive(Debug)]
pub struct GetTimeout(Option<String>);

impl fmt::Display for GetTimeout {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        try!(fmt.write_str(self.description()));
        if let Some(ref err) = self.0 {
            try!(write!(fmt, ": {}", err));
        }
        Ok(())
    }
}

impl Error for GetTimeout {
    fn description(&self) -> &str {
        "Timed out while waiting for a connection"
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<M>
    where M: ManageConnection
{
    pool: Pool<M>,
    conn: Option<Conn<M::Connection>>,
}

impl<M> fmt::Debug for PooledConnection<M>
    where M: ManageConnection,
          M::Connection: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.conn.as_ref().unwrap().conn, fmt)
    }
}

impl<M> Drop for PooledConnection<M>
    where M: ManageConnection
{
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<M> Deref for PooledConnection<M>
    where M: ManageConnection
{
    type Target = M::Connection;

    fn deref(&self) -> &M::Connection {
        &self.conn.as_ref().unwrap().conn
    }
}

impl<M> DerefMut for PooledConnection<M>
    where M: ManageConnection
{
    fn deref_mut(&mut self) -> &mut M::Connection {
        &mut self.conn.as_mut().unwrap().conn
    }
}
<commit_msg>Don't store idle_start in active connections<commit_after>//! A generic connection pool.
//!
//! Opening a new database connection every time one is needed is both
//! inefficient and can lead to resource exhaustion under high traffic
//! conditions. A connection pool maintains a set of open connections to a
//! database, handing them out for repeated use.
//!
//! r2d2 is agnostic to the connection type it is managing. Implementors of the
//! `ManageConnection` trait provide the database-specific logic to create and
//! check the health of connections.
//!
//! # Example
//!
//! Using an imaginary "foodb" database.
//!
//! ```rust,ignore
//! use std::thread;
//!
//! extern crate r2d2;
//! extern crate r2d2_foodb;
//!
//! fn main() {
//!     let config = r2d2::Config::builder()
//!         .pool_size(15)
//!         .build();
//!     let manager = r2d2_foodb::FooConnectionManager::new("localhost:1234");
//!
//!     let pool = r2d2::Pool::new(config, manager).unwrap();
//!
//!     for _ in 0..20 {
//!         let pool = pool.clone();
//!         thread::spawn(move || {
//!             let conn = pool.get().unwrap();
//!             // use the connection
//!             // it will be returned to the pool when it falls out of scope.
//!         })
//!     }
//! }
//! ```
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/r2d2/doc/v0.7.0")]

extern crate antidote;
#[macro_use]
extern crate log;

use antidote::{Mutex, Condvar};
use std::cmp;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::mem;
use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};

#[doc(inline)]
pub use config::Config;

use task::ScheduledThreadPool;

pub mod config;
mod task;
mod thunk;

#[cfg(test)]
mod test;

/// A trait which provides connection-specific functionality.
pub trait ManageConnection: Send + Sync + 'static {
    /// The connection type this manager deals with.
    type Connection: Send + 'static;

    /// The error type returned by `Connection`s.
    type Error: Error + 'static;

    /// Attempts to create a new connection.
    fn connect(&self) -> Result<Self::Connection, Self::Error>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut Self::Connection) -> bool;
}

/// A trait which handles errors reported by the `ManageConnection`.
pub trait HandleError<E>: fmt::Debug + Send + Sync + 'static {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

/// A `HandleError` implementation which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NopErrorHandler;

impl<E> HandleError<E> for NopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// A `HandleError` implementation which logs at the error level.
#[derive(Copy, Clone, Debug)]
pub struct LoggingErrorHandler;

impl<E> HandleError<E> for LoggingErrorHandler
    where E: Error
{
    fn handle_error(&self, error: E) {
        error!("{}", error);
    }
}

/// A trait which allows for customization of connections.
pub trait CustomizeConnection<C, E>: fmt::Debug + Send + Sync + 'static {
    /// Called with connections immediately after they are returned from
    /// `ManageConnection::connect`.
    ///
    /// The default implementation simply returns `Ok(())`.
    ///
    /// # Errors
    ///
    /// If this method returns an error, the connection will be discarded.
    #[allow(unused_variables)]
    fn on_acquire(&self, conn: &mut C) -> Result<(), E> {
        Ok(())
    }
}

/// A `CustomizeConnection` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NopConnectionCustomizer;

impl<C, E> CustomizeConnection<C, E> for NopConnectionCustomizer {}

struct Conn<C> {
    conn: C,
    birth: Instant,
}

struct IdleConn<C> {
    conn: Conn<C>,
    idle_start: Instant,
}

struct PoolInternals<C> {
    conns: VecDeque<IdleConn<C>>,
    num_conns: u32,
    pending_conns: u32,
    last_error: Option<String>,
}

struct SharedPool<M>
    where M: ManageConnection
{
    config: Config<M::Connection, M::Error>,
    manager: M,
    internals: Mutex<PoolInternals<M::Connection>>,
    cond: Condvar,
    thread_pool: ScheduledThreadPool,
}

fn drop_conn<M>(shared: &Arc<SharedPool<M>>, internals: &mut PoolInternals<M::Connection>)
    where M: ManageConnection
{
    internals.num_conns -= 1;

    let min = shared.config.min_idle().unwrap_or(shared.config.pool_size());
    if internals.num_conns + internals.pending_conns < min {
        add_connection(shared, internals);
    }
}

fn add_connection<M>(shared: &Arc<SharedPool<M>>, internals: &mut PoolInternals<M::Connection>)
    where M: ManageConnection
{
    internals.pending_conns += 1;
    inner(Duration::from_secs(0), shared);

    fn inner<M>(delay: Duration, shared: &Arc<SharedPool<M>>)
        where M: ManageConnection
    {
        let new_shared = Arc::downgrade(shared);
        shared.thread_pool.run_after(delay, move || {
            let shared = match new_shared.upgrade() {
                Some(shared) => shared,
                None => return,
            };

            let conn = shared.manager.connect().and_then(|mut conn| {
                shared.config.connection_customizer().on_acquire(&mut conn).map(|_| conn)
            });
            match conn {
                Ok(conn) => {
                    let mut internals = shared.internals.lock();
                    internals.last_error = None;
                    let now = Instant::now();
                    let conn = IdleConn {
                        conn: Conn {
                            conn: conn,
                            birth: now,
                        },
                        idle_start: now,
                    };
                    internals.conns.push_back(conn);
                    internals.pending_conns -= 1;
                    internals.num_conns += 1;
                    shared.cond.notify_one();
                }
                Err(err) => {
                    shared.internals.lock().last_error = Some(err.to_string());
                    shared.config.error_handler().handle_error(err);
                    let delay = cmp::max(Duration::from_millis(200), delay);
                    let delay = cmp::min(shared.config.connection_timeout() / 2, delay * 2);
                    inner(delay, &shared);
                }
            }
        });
    }
}

fn reap_connections<M>(shared: &Weak<SharedPool<M>>)
    where M: ManageConnection
{
    let shared = match shared.upgrade() {
        Some(shared) => shared,
        None => return,
    };

    let mut old = VecDeque::with_capacity(shared.config.pool_size() as usize);
    let mut to_drop = vec![];

    let mut internals = shared.internals.lock();
    mem::swap(&mut old, &mut internals.conns);
    let now = Instant::now();
    for conn in old {
        let mut reap = false;
        if let Some(timeout) = shared.config.idle_timeout() {
            reap |= now - conn.idle_start >= timeout;
        }
        if let Some(lifetime) = shared.config.max_lifetime() {
            reap |= now - conn.conn.birth >= lifetime;
        }
        if reap {
            drop_conn(&shared, &mut internals);
            to_drop.push(conn.conn.conn);
        } else {
            internals.conns.push_back(conn);
        }
    }
    drop(internals); // make sure we run to_drop destructors without this locked
}

/// A generic connection pool.
pub struct Pool<M: ManageConnection>(Arc<SharedPool<M>>);

/// Returns a new `Pool` referencing the same state as `self`.
impl<M> Clone for Pool<M>
    where M: ManageConnection
{
    fn clone(&self) -> Pool<M> {
        Pool(self.0.clone())
    }
}

impl<M> fmt::Debug for Pool<M>
    where M: ManageConnection + fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.0.internals.lock();

        fmt.debug_struct("Pool")
           .field("connections", &inner.num_conns)
           .field("idle_connections", &inner.conns.len())
           .field("config", &self.0.config)
           .field("manager", &self.0.manager)
           .finish()
    }
}

impl<M> Pool<M>
    where M: ManageConnection
{
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value if `initialization_fail_fast` is set to true in
    /// the configuration and the pool is unable to open all of its
    /// connections.
    pub fn new(config: Config<M::Connection, M::Error>,
               manager: M)
               -> Result<Pool<M>, InitializationError> {
        Pool::new_inner(config, manager, Duration::from_secs(30))
    }

    // for testing
    fn new_inner(config: Config<M::Connection, M::Error>,
                 manager: M,
                 reaper_rate: Duration)
                 -> Result<Pool<M>, InitializationError> {
        let internals = PoolInternals {
            conns: VecDeque::with_capacity(config.pool_size() as usize),
            num_conns: 0,
            pending_conns: 0,
            last_error: None,
        };

        let shared = Arc::new(SharedPool {
            thread_pool: ScheduledThreadPool::new(config.helper_threads() as usize),
            config: config,
            manager: manager,
            internals: Mutex::new(internals),
            cond: Condvar::new(),
        });

        let initial_size = shared.config.min_idle().unwrap_or(shared.config.pool_size());
        {
            let mut inner = shared.internals.lock();
            for _ in 0..initial_size {
                add_connection(&shared, &mut inner);
            }
            drop(inner);
        }

        if shared.config.initialization_fail_fast() {
            let end = Instant::now() + shared.config.connection_timeout();
            let mut internals = shared.internals.lock();

            while internals.num_conns != initial_size {
                let now = Instant::now();
                if now >= end {
                    return Err(InitializationError(internals.last_error.take()));
                }
                internals = shared.cond
                                  .wait_timeout(internals, end - now)
                                  .0;
            }
        }

        if shared.config.max_lifetime().is_some() || shared.config.idle_timeout().is_some() {
            let s = Arc::downgrade(&shared);
            shared.thread_pool
                  .run_at_fixed_rate(reaper_rate, move || reap_connections(&s));
        }

        Ok(Pool(shared))
    }

    /// Retrieves a connection from the pool.
    ///
    /// Waits for at most `Config::connection_timeout` before returning an
    /// error.
    pub fn get(&self) -> Result<PooledConnection<M>, GetTimeout> {
        let end = Instant::now() + self.0.config.connection_timeout();
        let mut internals = self.0.internals.lock();

        let connection;
        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.0.config.test_on_check_out() {
                        if let Err(e) = self.0.manager.is_valid(&mut conn.conn.conn) {
                            let msg = e.to_string();
                            self.0.config.error_handler().handle_error(e);
                            internals = self.0.internals.lock();
                            internals.last_error = Some(msg);
                            drop_conn(&self.0, &mut internals);
                            continue;
                        }
                    }

                    connection = conn.conn;
                    break;
                }
                None => {
                    if internals.num_conns + internals.pending_conns < self.0.config.pool_size() {
                        add_connection(&self.0, &mut internals);
                    }

                    let now = Instant::now();
                    if now >= end {
                        return Err(GetTimeout(internals.last_error.take()));
                    };
                    internals = self.0
                                    .cond
                                    .wait_timeout(internals, end - now)
                                    .0;
                }
            }
        }

        Ok(PooledConnection {
            pool: self.clone(),
            conn: Some(connection),
        })
    }

    fn put_back(&self, mut conn: Conn<M::Connection>) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.0.manager.has_broken(&mut conn.conn);

        let mut internals = self.0.internals.lock();
        if broken {
            drop_conn(&self.0, &mut internals);
        } else {
            let conn = IdleConn {
                conn: conn,
                idle_start: Instant::now(),
            };
            internals.conns.push_back(conn);
            self.0.cond.notify_one();
        }
    }
}

/// An error returned by `Pool::new` if it fails to initialize connections.
#[derive(Debug)]
pub struct InitializationError(Option<String>);

impl fmt::Display for InitializationError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        try!(fmt.write_str(self.description()));
        if let Some(ref err) = self.0 {
            try!(write!(fmt, ": {}", err));
        }
        Ok(())
    }
}

impl Error for InitializationError {
    fn description(&self) -> &str {
        "Unable to initialize connections"
    }
}

/// An error returned by `Pool::get` if it times out without retrieving a connection.
#[derive(Debug)]
pub struct GetTimeout(Option<String>);

impl fmt::Display for GetTimeout {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        try!(fmt.write_str(self.description()));
        if let Some(ref err) = self.0 {
            try!(write!(fmt, ": {}", err));
        }
        Ok(())
    }
}

impl Error for GetTimeout {
    fn description(&self) -> &str {
        "Timed out while waiting for a connection"
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<M>
    where M: ManageConnection
{
    pool: Pool<M>,
    conn: Option<Conn<M::Connection>>,
}

impl<M> fmt::Debug for PooledConnection<M>
    where M: ManageConnection,
          M::Connection: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.conn.as_ref().unwrap().conn, fmt)
    }
}

impl<M> Drop for PooledConnection<M>
    where M: ManageConnection
{
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<M> Deref for PooledConnection<M>
    where M: ManageConnection
{
    type Target = M::Connection;

    fn deref(&self) -> &M::Connection {
        &self.conn.as_ref().unwrap().conn
    }
}

impl<M> DerefMut for PooledConnection<M>
    where M: ManageConnection
{
    fn deref_mut(&mut self) -> &mut M::Connection {
        &mut self.conn.as_mut().unwrap().conn
    }
}
<|endoftext|>