<commit_before>pub mod bookitem;
pub mod bookconfig;

pub mod bookconfig_test;

pub use self::bookitem::{BookItem, BookItems};
pub use self::bookconfig::BookConfig;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::error::Error;
use std::io;
use std::io::Write;
use std::io::ErrorKind;
use std::process::Command;

use {theme, parse, utils};
use renderer::{Renderer, HtmlHandlebars};

pub struct MDBook {
    root: PathBuf,
    dest: PathBuf,
    src: PathBuf,
    theme_path: PathBuf,

    pub title: String,
    pub author: String,
    pub description: String,

    pub content: Vec<BookItem>,
    renderer: Box<Renderer>,

    livereload: Option<String>,

    /// Should `mdbook build` create files referenced from SUMMARY.md if they
    /// don't exist
    pub create_missing: bool,
}

impl MDBook {
    /// Create a new `MDBook` struct with root directory `root`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use std::path::Path;
    /// # fn main() {
    /// let book = MDBook::new(Path::new("root_dir"));
    /// # }
    /// ```
    ///
    /// In this example, `root_dir` will be the root directory of our book and is specified in function
    /// of the current working directory by using a relative path instead of an absolute path.
    ///
    /// Default directory paths:
    ///
    /// - source: `root/src`
    /// - output: `root/book`
    /// - theme: `root/theme`
    ///
    /// They can both be changed by using [`set_src()`](#method.set_src) and [`set_dest()`](#method.set_dest)

    pub fn new(root: &Path) -> MDBook {

        if !root.exists() || !root.is_dir() {
            warn!("{:?} No directory with that name", root);
        }

        MDBook {
            root: root.to_owned(),
            dest: root.join("book"),
            src: root.join("src"),
            theme_path: root.join("theme"),

            title: String::new(),
            author: String::new(),
            description: String::new(),

            content: vec![],
            renderer: Box::new(HtmlHandlebars::new()),

            livereload: None,
            create_missing: true,
        }
    }

    /// Returns a flat depth-first iterator over the elements of the book, it returns an [BookItem enum](bookitem.html):
    /// `(section: String, bookitem: &BookItem)`
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use mdbook::BookItem;
    /// # use std::path::Path;
    /// # fn main() {
    /// # let mut book = MDBook::new(Path::new("mybook"));
    /// for item in book.iter() {
    ///     match item {
    ///         &BookItem::Chapter(ref section, ref chapter) => {},
    ///         &BookItem::Affix(ref chapter) => {},
    ///         &BookItem::Spacer => {},
    ///     }
    /// }
    ///
    /// // would print something like this:
    /// // 1. Chapter 1
    /// // 1.1 Sub Chapter
    /// // 1.2 Sub Chapter
    /// // 2. Chapter 2
    /// //
    /// // etc.
    /// # }
    /// ```

    pub fn iter(&self) -> BookItems {
        BookItems {
            items: &self.content[..],
            current_index: 0,
            stack: Vec::new(),
        }
    }

    /// `init()` creates some boilerplate files and directories to get you started with your book.
    ///
    /// ```text
    /// book-test/
    /// ├── book
    /// └── src
    ///     ├── chapter_1.md
    ///     └── SUMMARY.md
    /// ```
    ///
    /// It uses the paths given as source and output directories and adds a `SUMMARY.md` and a
    /// `chapter_1.md` to the source directory.

    pub fn init(&mut self) -> Result<(), Box<Error>> {

        debug!("[fn]: init");

        if !self.root.exists() {
            fs::create_dir_all(&self.root).unwrap();
            info!("{:?} created", &self.root);
        }

        {

            if !self.dest.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.dest);
                try!(fs::create_dir_all(&self.dest));
            }

            if !self.src.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.src);
                try!(fs::create_dir_all(&self.src));
            }

            let summary = self.src.join("SUMMARY.md");

            if !summary.exists() {

                // Summary does not exist, create it

                debug!("[*]: {:?} does not exist, trying to create SUMMARY.md", self.src.join("SUMMARY.md"));
                let mut f = try!(File::create(&self.src.join("SUMMARY.md")));

                debug!("[*]: Writing to SUMMARY.md");

                try!(writeln!(f, "# Summary"));
                try!(writeln!(f, ""));
                try!(writeln!(f, "- [Chapter 1](./chapter_1.md)"));
            }
        }

        // parse SUMMARY.md, and create the missing item related file
        try!(self.parse_summary());

        debug!("[*]: constructing paths for missing files");
        for item in self.iter() {
            debug!("[*]: item: {:?}", item);
            match *item {
                BookItem::Spacer => continue,
                BookItem::Chapter(_, ref ch) |
                BookItem::Affix(ref ch) => {
                    if ch.path != PathBuf::new() {
                        let path = self.src.join(&ch.path);

                        if !path.exists() {
                            if !self.create_missing {
                                return Err(format!("'{}' referenced from SUMMARY.md does not exist.", path.to_string_lossy()).into());
                            }
                            debug!("[*]: {:?} does not exist, trying to create file", path);
                            try!(::std::fs::create_dir_all(path.parent().unwrap()));
                            let mut f = try!(File::create(path));

                            // debug!("[*]: Writing to {:?}", path);
                            try!(writeln!(f, "# {}", ch.name));
                        }
                    }
                },
            }
        }

        debug!("[*]: init done");
        Ok(())
    }

    pub fn create_gitignore(&self) {
        let gitignore = self.get_gitignore();

        if !gitignore.exists() {
            // Gitignore does not exist, create it

            // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`. If it
            // is not, `strip_prefix` will return an Error.
            if !self.get_dest().starts_with(&self.root) {
                return;
            }

            let relative = self.get_dest()
                               .strip_prefix(&self.root)
                               .expect("Destination is not relative to root.");
            let relative = relative.to_str()
                                   .expect("Path could not be yielded into a string slice.");

            debug!("[*]: {:?} does not exist, trying to create .gitignore", gitignore);

            let mut f = File::create(&gitignore).expect("Could not create file.");

            debug!("[*]: Writing to .gitignore");

            writeln!(f, "{}", relative).expect("Could not write to file.");
        }
    }

    /// The `build()` method is the one where everything happens. First it parses `SUMMARY.md` to
    /// construct the book's structure in the form of a `Vec<BookItem>` and then calls `render()`
    /// method of the current renderer.
    ///
    /// It is the renderer who generates all the output files.
    pub fn build(&mut self) -> Result<(), Box<Error>> {
        debug!("[fn]: build");

        try!(self.init());

        // Clean output directory
        try!(utils::fs::remove_dir_content(&self.dest));

        try!(self.renderer.render(&self));

        Ok(())
    }

    pub fn get_gitignore(&self) -> PathBuf {
        self.root.join(".gitignore")
    }

    pub fn copy_theme(&self) -> Result<(), Box<Error>> {
        debug!("[fn]: copy_theme");

        let theme_dir = self.src.join("theme");

        if !theme_dir.exists() {
            debug!("[*]: {:?} does not exist, trying to create directory", theme_dir);
            try!(fs::create_dir(&theme_dir));
        }

        // index.hbs
        let mut index = try!(File::create(&theme_dir.join("index.hbs")));
        try!(index.write_all(theme::INDEX));

        // book.css
        let mut css = try!(File::create(&theme_dir.join("book.css")));
        try!(css.write_all(theme::CSS));

        // favicon.png
        let mut favicon = try!(File::create(&theme_dir.join("favicon.png")));
        try!(favicon.write_all(theme::FAVICON));

        // book.js
        let mut js = try!(File::create(&theme_dir.join("book.js")));
        try!(js.write_all(theme::JS));

        // highlight.css
        let mut highlight_css = try!(File::create(&theme_dir.join("highlight.css")));
        try!(highlight_css.write_all(theme::HIGHLIGHT_CSS));

        // highlight.js
        let mut highlight_js = try!(File::create(&theme_dir.join("highlight.js")));
        try!(highlight_js.write_all(theme::HIGHLIGHT_JS));

        Ok(())
    }

    pub fn write_file<P: AsRef<Path>>(&self, filename: P, content: &[u8]) -> Result<(), Box<Error>> {
        let path = self.get_dest().join(filename);
        try!(utils::fs::create_file(&path).and_then(|mut file| {
            file.write_all(content)
        }).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("Could not create {}: {}", path.display(), e))
        }));
        Ok(())
    }

    /// Parses the `book.json` file (if it exists) to extract the configuration parameters.
    /// The `book.json` file should be in the root directory of the book.
    /// The root directory is the one specified when creating a new `MDBook`

    pub fn read_config(mut self) -> Self {

        let config = BookConfig::new(&self.root)
                         .read_config(&self.root)
                         .to_owned();

        self.title = config.title;
        self.description = config.description;
        self.author = config.author;

        self.dest = config.dest;
        self.src = config.src;
        self.theme_path = config.theme_path;

        self
    }

    /// You can change the default renderer to another one by using this method. The only requirement
    /// is for your renderer to implement the [Renderer trait](../../renderer/renderer/trait.Renderer.html)
    ///
    /// ```no_run
    /// extern crate mdbook;
    /// use mdbook::MDBook;
    /// use mdbook::renderer::HtmlHandlebars;
    /// # use std::path::Path;
    ///
    /// fn main() {
    ///     let mut book = MDBook::new(Path::new("mybook"))
    ///                         .set_renderer(Box::new(HtmlHandlebars::new()));
    ///
    ///     // In this example we replace the default renderer by the default renderer...
    ///     // Don't forget to put your renderer in a Box
    /// }
    /// ```
    ///
    /// **note:** Don't forget to put your renderer in a `Box` before passing it to `set_renderer()`

    pub fn set_renderer(mut self, renderer: Box<Renderer>) -> Self {
        self.renderer = renderer;
        self
    }

    pub fn test(&mut self) -> Result<(), Box<Error>> {
        // read in the chapters
        try!(self.parse_summary());
        for item in self.iter() {

            if let BookItem::Chapter(_, ref ch) = *item {
                if ch.path != PathBuf::new() {

                    let path = self.get_src().join(&ch.path);

                    println!("[*]: Testing file: {:?}", path);

                    let output_result = Command::new("rustdoc")
                                            .arg(&path)
                                            .arg("--test")
                                            .output();
                    let output = try!(output_result);

                    if !output.status.success() {
                        return Err(Box::new(io::Error::new(ErrorKind::Other, format!(
                                        "{}\n{}",
                                        String::from_utf8_lossy(&output.stdout),
                                        String::from_utf8_lossy(&output.stderr)))) as Box<Error>);
                    }
                }
            }
        }
        Ok(())
    }

    pub fn get_root(&self) -> &Path {
        &self.root
    }

    pub fn set_dest(mut self, dest: &Path) -> Self {

        // Handle absolute and relative paths
        if dest.is_absolute() {
            self.dest = dest.to_owned();
        } else {
            let dest = self.root.join(dest).to_owned();
            self.dest = dest;
        }

        self
    }

    pub fn get_dest(&self) -> &Path {
        &self.dest
    }

    pub fn set_src(mut self, src: &Path) -> Self {

        // Handle absolute and relative paths
        if src.is_absolute() {
            self.src = src.to_owned();
        } else {
            let src = self.root.join(src).to_owned();
            self.src = src;
        }

        self
    }

    pub fn get_src(&self) -> &Path {
        &self.src
    }

    pub fn set_title(mut self, title: &str) -> Self {
        self.title = title.to_owned();
        self
    }

    pub fn get_title(&self) -> &str {
        &self.title
    }

    pub fn set_author(mut self, author: &str) -> Self {
        self.author = author.to_owned();
        self
    }

    pub fn get_author(&self) -> &str {
        &self.author
    }

    pub fn set_description(mut self, description: &str) -> Self {
        self.description = description.to_owned();
        self
    }

    pub fn get_description(&self) -> &str {
        &self.description
    }

    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
        self.livereload = Some(livereload);
        self
    }

    pub fn unset_livereload(&mut self) -> &Self {
        self.livereload = None;
        self
    }

    pub fn get_livereload(&self) -> Option<&String> {
        self.livereload.as_ref()
    }

    pub fn set_theme_path(mut self, theme_path: &Path) -> Self {
        self.theme_path = if theme_path.is_absolute() {
            theme_path.to_owned()
        } else {
            self.root.join(theme_path).to_owned()
        };
        self
    }

    pub fn get_theme_path(&self) -> &Path {
        &self.theme_path
    }

    // Construct book
    fn parse_summary(&mut self) -> Result<(), Box<Error>> {
        // When append becomes stable, use self.content.append() ...
        self.content = try!(parse::construct_bookitems(&self.src.join("SUMMARY.md")));
        Ok(())
    }
}
<commit_msg>Refactor to prevent excessive indentation.<commit_after>pub mod bookitem;
pub mod bookconfig;

pub mod bookconfig_test;

pub use self::bookitem::{BookItem, BookItems};
pub use self::bookconfig::BookConfig;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::error::Error;
use std::io;
use std::io::Write;
use std::io::ErrorKind;
use std::process::Command;

use {theme, parse, utils};
use renderer::{Renderer, HtmlHandlebars};

pub struct MDBook {
    root: PathBuf,
    dest: PathBuf,
    src: PathBuf,
    theme_path: PathBuf,

    pub title: String,
    pub author: String,
    pub description: String,

    pub content: Vec<BookItem>,
    renderer: Box<Renderer>,

    livereload: Option<String>,

    /// Should `mdbook build` create files referenced from SUMMARY.md if they
    /// don't exist
    pub create_missing: bool,
}

impl MDBook {
    /// Create a new `MDBook` struct with root directory `root`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use std::path::Path;
    /// # fn main() {
    /// let book = MDBook::new(Path::new("root_dir"));
    /// # }
    /// ```
    ///
    /// In this example, `root_dir` will be the root directory of our book and is specified in function
    /// of the current working directory by using a relative path instead of an absolute path.
    ///
    /// Default directory paths:
    ///
    /// - source: `root/src`
    /// - output: `root/book`
    /// - theme: `root/theme`
    ///
    /// They can both be changed by using [`set_src()`](#method.set_src) and [`set_dest()`](#method.set_dest)

    pub fn new(root: &Path) -> MDBook {

        if !root.exists() || !root.is_dir() {
            warn!("{:?} No directory with that name", root);
        }

        MDBook {
            root: root.to_owned(),
            dest: root.join("book"),
            src: root.join("src"),
            theme_path: root.join("theme"),

            title: String::new(),
            author: String::new(),
            description: String::new(),

            content: vec![],
            renderer: Box::new(HtmlHandlebars::new()),

            livereload: None,
            create_missing: true,
        }
    }

    /// Returns a flat depth-first iterator over the elements of the book, it returns an [BookItem enum](bookitem.html):
    /// `(section: String, bookitem: &BookItem)`
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use mdbook::BookItem;
    /// # use std::path::Path;
    /// # fn main() {
    /// # let mut book = MDBook::new(Path::new("mybook"));
    /// for item in book.iter() {
    ///     match item {
    ///         &BookItem::Chapter(ref section, ref chapter) => {},
    ///         &BookItem::Affix(ref chapter) => {},
    ///         &BookItem::Spacer => {},
    ///     }
    /// }
    ///
    /// // would print something like this:
    /// // 1. Chapter 1
    /// // 1.1 Sub Chapter
    /// // 1.2 Sub Chapter
    /// // 2. Chapter 2
    /// //
    /// // etc.
    /// # }
    /// ```

    pub fn iter(&self) -> BookItems {
        BookItems {
            items: &self.content[..],
            current_index: 0,
            stack: Vec::new(),
        }
    }

    /// `init()` creates some boilerplate files and directories to get you started with your book.
    ///
    /// ```text
    /// book-test/
    /// ├── book
    /// └── src
    ///     ├── chapter_1.md
    ///     └── SUMMARY.md
    /// ```
    ///
    /// It uses the paths given as source and output directories and adds a `SUMMARY.md` and a
    /// `chapter_1.md` to the source directory.

    pub fn init(&mut self) -> Result<(), Box<Error>> {

        debug!("[fn]: init");

        if !self.root.exists() {
            fs::create_dir_all(&self.root).unwrap();
            info!("{:?} created", &self.root);
        }

        {

            if !self.dest.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.dest);
                try!(fs::create_dir_all(&self.dest));
            }

            if !self.src.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.src);
                try!(fs::create_dir_all(&self.src));
            }

            let summary = self.src.join("SUMMARY.md");

            if !summary.exists() {

                // Summary does not exist, create it

                debug!("[*]: {:?} does not exist, trying to create SUMMARY.md", self.src.join("SUMMARY.md"));
                let mut f = try!(File::create(&self.src.join("SUMMARY.md")));

                debug!("[*]: Writing to SUMMARY.md");

                try!(writeln!(f, "# Summary"));
                try!(writeln!(f, ""));
                try!(writeln!(f, "- [Chapter 1](./chapter_1.md)"));
            }
        }

        // parse SUMMARY.md, and create the missing item related file
        try!(self.parse_summary());

        debug!("[*]: constructing paths for missing files");
        for item in self.iter() {
            debug!("[*]: item: {:?}", item);
            let ch = match *item {
                BookItem::Spacer => continue,
                BookItem::Chapter(_, ref ch) |
                BookItem::Affix(ref ch) => ch,
            };
            if ch.path != PathBuf::new() {
                let path = self.src.join(&ch.path);

                if !path.exists() {
                    if !self.create_missing {
                        return Err(format!("'{}' referenced from SUMMARY.md does not exist.", path.to_string_lossy()).into());
                    }
                    debug!("[*]: {:?} does not exist, trying to create file", path);
                    try!(::std::fs::create_dir_all(path.parent().unwrap()));
                    let mut f = try!(File::create(path));

                    // debug!("[*]: Writing to {:?}", path);
                    try!(writeln!(f, "# {}", ch.name));
                }
            }
        }

        debug!("[*]: init done");
        Ok(())
    }

    pub fn create_gitignore(&self) {
        let gitignore = self.get_gitignore();

        if !gitignore.exists() {
            // Gitignore does not exist, create it

            // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`. If it
            // is not, `strip_prefix` will return an Error.
            if !self.get_dest().starts_with(&self.root) {
                return;
            }

            let relative = self.get_dest()
                               .strip_prefix(&self.root)
                               .expect("Destination is not relative to root.");
            let relative = relative.to_str()
                                   .expect("Path could not be yielded into a string slice.");

            debug!("[*]: {:?} does not exist, trying to create .gitignore", gitignore);

            let mut f = File::create(&gitignore).expect("Could not create file.");

            debug!("[*]: Writing to .gitignore");

            writeln!(f, "{}", relative).expect("Could not write to file.");
        }
    }

    /// The `build()` method is the one where everything happens. First it parses `SUMMARY.md` to
    /// construct the book's structure in the form of a `Vec<BookItem>` and then calls `render()`
    /// method of the current renderer.
    ///
    /// It is the renderer who generates all the output files.
    pub fn build(&mut self) -> Result<(), Box<Error>> {
        debug!("[fn]: build");

        try!(self.init());

        // Clean output directory
        try!(utils::fs::remove_dir_content(&self.dest));

        try!(self.renderer.render(&self));

        Ok(())
    }

    pub fn get_gitignore(&self) -> PathBuf {
        self.root.join(".gitignore")
    }

    pub fn copy_theme(&self) -> Result<(), Box<Error>> {
        debug!("[fn]: copy_theme");

        let theme_dir = self.src.join("theme");

        if !theme_dir.exists() {
            debug!("[*]: {:?} does not exist, trying to create directory", theme_dir);
            try!(fs::create_dir(&theme_dir));
        }

        // index.hbs
        let mut index = try!(File::create(&theme_dir.join("index.hbs")));
        try!(index.write_all(theme::INDEX));

        // book.css
        let mut css = try!(File::create(&theme_dir.join("book.css")));
        try!(css.write_all(theme::CSS));

        // favicon.png
        let mut favicon = try!(File::create(&theme_dir.join("favicon.png")));
        try!(favicon.write_all(theme::FAVICON));

        // book.js
        let mut js = try!(File::create(&theme_dir.join("book.js")));
        try!(js.write_all(theme::JS));

        // highlight.css
        let mut highlight_css = try!(File::create(&theme_dir.join("highlight.css")));
        try!(highlight_css.write_all(theme::HIGHLIGHT_CSS));

        // highlight.js
        let mut highlight_js = try!(File::create(&theme_dir.join("highlight.js")));
        try!(highlight_js.write_all(theme::HIGHLIGHT_JS));

        Ok(())
    }

    pub fn write_file<P: AsRef<Path>>(&self, filename: P, content: &[u8]) -> Result<(), Box<Error>> {
        let path = self.get_dest().join(filename);
        try!(utils::fs::create_file(&path).and_then(|mut file| {
            file.write_all(content)
        }).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("Could not create {}: {}", path.display(), e))
        }));
        Ok(())
    }

    /// Parses the `book.json` file (if it exists) to extract the configuration parameters.
    /// The `book.json` file should be in the root directory of the book.
    /// The root directory is the one specified when creating a new `MDBook`

    pub fn read_config(mut self) -> Self {

        let config = BookConfig::new(&self.root)
                         .read_config(&self.root)
                         .to_owned();

        self.title = config.title;
        self.description = config.description;
        self.author = config.author;

        self.dest = config.dest;
        self.src = config.src;
        self.theme_path = config.theme_path;

        self
    }

    /// You can change the default renderer to another one by using this method. The only requirement
    /// is for your renderer to implement the [Renderer trait](../../renderer/renderer/trait.Renderer.html)
    ///
    /// ```no_run
    /// extern crate mdbook;
    /// use mdbook::MDBook;
    /// use mdbook::renderer::HtmlHandlebars;
    /// # use std::path::Path;
    ///
    /// fn main() {
    ///     let mut book = MDBook::new(Path::new("mybook"))
    ///                         .set_renderer(Box::new(HtmlHandlebars::new()));
    ///
    ///     // In this example we replace the default renderer by the default renderer...
    ///     // Don't forget to put your renderer in a Box
    /// }
    /// ```
    ///
    /// **note:** Don't forget to put your renderer in a `Box` before passing it to `set_renderer()`

    pub fn set_renderer(mut self, renderer: Box<Renderer>) -> Self {
        self.renderer = renderer;
        self
    }

    pub fn test(&mut self) -> Result<(), Box<Error>> {
        // read in the chapters
        try!(self.parse_summary());
        for item in self.iter() {

            if let BookItem::Chapter(_, ref ch) = *item {
                if ch.path != PathBuf::new() {

                    let path = self.get_src().join(&ch.path);

                    println!("[*]: Testing file: {:?}", path);

                    let output_result = Command::new("rustdoc")
                                            .arg(&path)
                                            .arg("--test")
                                            .output();
                    let output = try!(output_result);

                    if !output.status.success() {
                        return Err(Box::new(io::Error::new(ErrorKind::Other, format!(
                                        "{}\n{}",
                                        String::from_utf8_lossy(&output.stdout),
                                        String::from_utf8_lossy(&output.stderr)))) as Box<Error>);
                    }
                }
            }
        }
        Ok(())
    }

    pub fn get_root(&self) -> &Path {
        &self.root
    }

    pub fn set_dest(mut self, dest: &Path) -> Self {

        // Handle absolute and relative paths
        if dest.is_absolute() {
            self.dest = dest.to_owned();
        } else {
            let dest = self.root.join(dest).to_owned();
            self.dest = dest;
        }

        self
    }

    pub fn get_dest(&self) -> &Path {
        &self.dest
    }

    pub fn set_src(mut self, src: &Path) -> Self {

        // Handle absolute and relative paths
        if src.is_absolute() {
            self.src = src.to_owned();
        } else {
            let src = self.root.join(src).to_owned();
            self.src = src;
        }

        self
    }

    pub fn get_src(&self) -> &Path {
        &self.src
    }

    pub fn set_title(mut self, title: &str) -> Self {
        self.title = title.to_owned();
        self
    }

    pub fn get_title(&self) -> &str {
        &self.title
    }

    pub fn set_author(mut self, author: &str) -> Self {
        self.author = author.to_owned();
        self
    }

    pub fn get_author(&self) -> &str {
        &self.author
    }

    pub fn set_description(mut self, description: &str) -> Self {
        self.description = description.to_owned();
        self
    }

    pub fn get_description(&self) -> &str {
        &self.description
    }

    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
        self.livereload = Some(livereload);
        self
    }

    pub fn unset_livereload(&mut self) -> &Self {
        self.livereload = None;
        self
    }

    pub fn get_livereload(&self) -> Option<&String> {
        self.livereload.as_ref()
    }

    pub fn set_theme_path(mut self, theme_path: &Path) -> Self {
        self.theme_path = if theme_path.is_absolute() {
            theme_path.to_owned()
        } else {
            self.root.join(theme_path).to_owned()
        };
        self
    }

    pub fn get_theme_path(&self) -> &Path {
        &self.theme_path
    }

    // Construct book
    fn parse_summary(&mut self) -> Result<(), Box<Error>> {
        // When append becomes stable, use self.content.append() ...
        self.content = try!(parse::construct_bookitems(&self.src.join("SUMMARY.md")));
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! A Rust snippet runner.
//!
//! Please see [readme](https://github.com/stevedonovan/runner/blob/master/readme.md)
extern crate easy_shortcuts as es;
extern crate lapp;
use es::traits::*;
use std::process;
use std::env;
use std::fs;
use std::path::PathBuf;
use std::collections::HashMap;
use std::io::Write;

mod crate_utils;
mod platform;

use platform::{open,edit,EXE,SO};

fn rustup_lib() -> String {
    es::shell("rustc --print sysroot") + "/lib"
}

// this will be initially written to ~/.cargo/.runner/prelude and
// can then be edited.
const PRELUDE: &'static str = "
#![allow(unused_imports)]
#![allow(dead_code)]
use std::fs;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::env;
use std::path::{PathBuf,Path};
#[allow(unused_macros)]
macro_rules! debug {
    ($x:expr) => {
        println!(\"{} = {:?}\",stringify!($x),$x);
    }
}

";

fn runner_directory() -> PathBuf {
    let mut home = crate_utils::cargo_home();
    home.push(".runner");
    home
}

fn cargo(args: &[&str]) {
    let res = process::Command::new("cargo")
        .args(args)
        .status()
        .or_die("can't run cargo");
    if ! res.success() {
        es::quit("cargo failed");
    }
}

const STATIC_CACHE: &str = "static-cache";
const DYNAMIC_CACHE: &str = "dy-cache";

fn static_cache_dir() -> PathBuf {
    runner_directory().join(STATIC_CACHE)
}

fn static_cache_dir_check(args: &lapp::Args) -> PathBuf {
    let static_cache = static_cache_dir();
    if ! static_cache.exists() {
        args.quit("please build static cache with --create first");
    }
    static_cache
}

fn build_static_cache() {
    cargo(&["build"]);
    cargo(&["build","--release"]);
    cargo(&["doc"]);
}

fn create_static_cache(crates: &[String], create: bool) {
    use std::io::prelude::*;
    let mut home = runner_directory();
    env::set_current_dir(&home).or_die("cannot change to home directory");
    if create {
        cargo(&["new","--bin",STATIC_CACHE]);
    }
    home.push(STATIC_CACHE);
    env::set_current_dir(&home).or_die("could not create static cache");
    {
        let mut deps = fs::OpenOptions::new().append(true)
            .open("Cargo.toml").or_die("could not append to cargo.toml");
        for c in crates {
            write!(deps,"{}=\"*\"\n",c).or_die("could not modify cargo.toml");
        }
    }
    build_static_cache();
}

// this is always called first and has the important role to ensure that
// runner's directory structure is created properly.
fn get_prelude() -> String {
    let home = runner_directory();
    let pristine = ! home.is_dir();
    if pristine {
        fs::create_dir(&home).or_die("cannot create runner directory");
    }
    let prelude = home.join("prelude");
    let bin = home.join("bin");
    if pristine {
        es::write_all(&prelude,PRELUDE);
        fs::create_dir(&home.join(DYNAMIC_CACHE)).or_die("cannot create dynamic cache");
    }
    if pristine || ! bin.is_dir() {
        fs::create_dir(&bin).or_die("cannot create output directory");
    }
    es::read_to_string(&prelude)
}

fn get_cache(build_static: bool, optimize: bool) -> PathBuf {
    let mut home = runner_directory();
    let scache;
    let cache = if build_static {
        let dir = if optimize {"release"} else {"debug"};
        scache = format!("{}/target/{}/deps",STATIC_CACHE,dir);
        &scache
    } else {DYNAMIC_CACHE};
    home.push(cache);
    home
}

fn add_aliases(aliases: Vec<String>) {
    if aliases.len() == 0 { return; }
    let alias_file = runner_directory().join("alias");
    let mut f = if alias_file.is_file() {
        fs::OpenOptions::new().append(true).open(&alias_file)
    } else {
        fs::File::create(&alias_file)
    }.or_die("cannot open runner alias file");

    for crate_alias in aliases {
        write!(f,"{}\n",crate_alias).or_die("cannot write to runner alias file");
    }
}

fn get_aliases() -> HashMap<String,String> {
    let alias_file = runner_directory().join("alias");
    if ! alias_file.is_file() { return HashMap::new(); }
	es::lines(es::open(&alias_file))
      .filter_map(|s| s.split_at_delim('=').trim()) // split into (String,String)
	  .to_map()
}

fn massage_snippet(code: String, prelude: String, extern_crates: Vec<String>, wild_crates: Vec<String>) -> String {
    fn indent_line(line: &str) -> String {
        format!("    {}\n",line)
    }
    let mut prefix = prelude;
    let mut body = String::new();
    {
        if extern_crates.len() > 0 {
            let aliases = get_aliases();
            for c in extern_crates {
                prefix += &if let Some(aliased) = aliases.get(&c) {
                    format!("extern crate {} as {};",aliased,c)
                } else {
                    format!("extern crate {};",c)
                };
            }
            for c in wild_crates {
                prefix += &format!("use {}::*;",c);
            }
        }
        let mut lines = code.lines();
        let mut first = true;
        for line in lines.by_ref() {
            let line = line.trim_left();
            if first { // files may start with #! shebang...
                if line.starts_with("#!") {
                    continue;
                }
                first = false;
            }
            if line.starts_with("//") || line.starts_with("#[macro_use") ||
                line.starts_with("extern ") || line.starts_with("use ") {
                prefix += line;
                prefix.push('\n');
            } else
            if line.len() > 0 {
                body += &indent_line(line);
                break;
            }
        }
        // and indent the rest!
        body.extend(lines.map(indent_line));
    }

    format!("{}
use std::error::Error;
fn run() -> Result<(),Box<Error>> {{
{}    Ok(())
}}
fn main() {{
    run().unwrap();
}}
        ",prefix,body)

}

const USAGE: &str = "
Compile and run small Rust snippets
  -s, --static build statically (default is dynamic)
  -O, --optimize optimized static build
  -e, --expression evaluate an expression
  -i, --iterator evaluate an iterator
  -n, --lines evaluate expression over stdin; 'line' is defined
  -x, --extern... (string) add an extern crate to the snippet
  -X, --wild... (string) like -x but implies wildcard use

  Cache Management:
  --create (string...) initialize the static cache with crates
  --add  (string...) add new crates to the cache (after --create)
  --edit  edit the static cache Cargo.toml
  --build rebuild the static cache
  --doc  display documentation
  --edit-prelude edit the default prelude for snippets
  --alias (string...) crate aliases in form alias=crate_name (used with -x)

  Dynamic compilation:
  -P, --crate-path show path of crate source in Cargo cache
  -C, --compile  compile crate dynamically (limited)
  --cfg... (string) pass configuration variables to rustc
  --libc  link dynamically against libc (special case)

  <program> (string) Rust program, snippet or expression
  <args> (string...) arguments to pass to program
";

fn main() {
    let args = lapp::parse_args(USAGE);
    let prelude = get_prelude();

    let aliases = args.get_strings("alias");
    if aliases.len() > 0 {
        add_aliases(aliases);
        return;
    }

    if args.get_bool("edit-prelude") {
        let rdir = runner_directory().join("prelude");
        edit(&rdir);
        return;
    }

    // Static Cache Management
    let crates = args.get_strings("create");
    if crates.len() > 0 {
        create_static_cache(&crates,true);
        return;
    }
    let crates = args.get_strings("add");
    if crates.len() > 0 {
        static_cache_dir_check(&args);
        create_static_cache(&crates,false);
        return;
    }

    if args.get_bool("edit") || args.get_bool("build") || args.get_bool("doc") {
        let static_cache = static_cache_dir_check(&args);
        if args.get_bool("build") {
            env::set_current_dir(&static_cache).or_die("static cache wasn't a directory?");
            build_static_cache();
        } else
        if args.get_bool("doc") {
            let docs = static_cache.join("target/doc/static_cache/index.html");
            open(&docs);
        } else {
            let toml = static_cache.join("Cargo.toml");
            edit(&toml);
        }
        return;
    }

    let first_arg = args.get_string("program");
    let file = PathBuf::from(&first_arg);

    // Dynamically linking crates (experimental!)
    if args.get_bool("crate-path") || args.get_bool("compile") {
        let (crate_path,crate_name) = if file.exists() {
            let filename = crate_utils::path_file_name(&file);
            if file.is_dir() { // assumed to be Cargo directory
                if ! file.join("Cargo.toml").exists() {
                    args.quit(&format!("not a Cargo project directory: {}",file.display()));
                }
                (file.join("src/lib.rs"), filename)
            } else { // should be just a Rust source file
                if file.extension().or_die("expecting extension") != "rs" {
                    args.quit("expecting Rust source file");
                }
                let name = crate_utils::path_file_name(&file.with_extension(""));
                (file, name)
            }
        } else {
            (crate_utils::cache_path(&first_arg).join("src/lib.rs"), first_arg)
        };
        if args.get_bool("crate-path") {
            println!("{}",crate_utils::cache_path(&crate_name).display());
        } else {
            let valid_crate_name = crate_name.replace('-',"_");
            let cache = get_cache(false, false);
            let mut builder = process::Command::new("rustc");
                builder.args(&["-C","prefer-dynamic"]).args(&["-C","debuginfo=0"])
                .arg("-L").arg(&cache)
                .args(&["--crate-type","dylib"])
                .arg("--out-dir").arg(&cache)
                .arg("--crate-name").arg(&valid_crate_name)
                .arg(crate_path);
           for c in args.get_strings("cfg") {
                builder.arg("--cfg").arg(&c);
           }
           if args.get_bool("libc") {
                builder.arg("--extern").arg(&format!("libc={}/liblibc.{}",cache.display(),SO));
           }
           builder.status().or_die("can't run rustc");
        }
        return;
    }

    let build_static = args.get_bool("static");
    let optimize = args.get_bool("optimize");

    // we'll pass rest of arguments to program
    let program_args = args.get_strings("args");

    let cache = get_cache(build_static, optimize);

    let mut snippet = false;
    let mut code = if args.get_bool("expression") {
        // Evaluating an expression: just print it out.
        format!("println!(\"{{:?}}\",{});", first_arg)
    } else
    if args.get_bool("iterator") {
        // The expression is anything that implements IntoIterator
        format!("let iter = {};\n for val in iter {{ println!(\"{{:?}}\",val);}}", first_arg)
    } else
    if args.get_bool("lines") {
        // The variable 'line' is available to an expression, evaluated for each line in stdin
        format!("
            let stdin = io::stdin();
            for line in stdin.lock().lines() {{
                let line = line?;
                let val = {};
                println!(\"{{:?}}\",val);
            }}
            ", first_arg)
    } else { // otherwise, just a file
        snippet = true;
        es::read_to_string(&file)
    };

    // proper Rust programs are accepted (this is a bit rough)
    if code.find("fn main").is_none() {
        let mut extern_crates = args.get_strings("extern");
        let wild_crates = args.get_strings("wild");
        if wild_crates.len() > 0 {
            extern_crates.extend(wild_crates.iter().cloned());
        }
        code = massage_snippet(code,prelude, extern_crates, wild_crates);
    }

    // we are going to put the expanded source and resulting exe in the runner bin dir
    let mut out_file = runner_directory().join("bin");
    if snippet {
        out_file.push(&file);
        out_file.set_extension("rs");
    } else {
        out_file.push("tmp.rs");
    }
    let mut program = out_file.clone();
    program.set_extension(EXE);

    es::write_all(&out_file,&code);

    let mut builder = process::Command::new("rustc");
    if ! build_static {
        builder.args(&["-C","prefer-dynamic"]).args(&["-C","debuginfo=0"]);
    } else {
        builder.arg(if optimize {"-O"} else {"-g"});
    }
    builder.arg("-L").arg(&cache);
    let status = builder.arg("-o").arg(&program)
        .arg(&out_file)
        .status().or_die("can't run rustc");
    if ! status.success() {
        return;
    }

    let mut builder = process::Command::new(&program);
    if ! build_static {
        builder.env("LD_LIBRARY_PATH",format!("{}:{}",rustup_lib(),cache.display()));
    }
    builder.args(&program_args)
        .status()
        .or_die(&format!("can't run program {:?}",program));

}

<commit_msg>--add,--create: recover from Cargo errors such as crate not found<commit_after>//! A Rust snippet runner.
//!
//! Please see [readme](https://github.com/stevedonovan/runner/blob/master/readme.md)
extern crate easy_shortcuts as es;
extern crate lapp;
use es::traits::*;
use std::process;
use std::env;
use std::fs;
use std::path::PathBuf;
use std::collections::HashMap;
use std::io::Write;

mod crate_utils;
mod platform;

use platform::{open,edit,EXE,SO};

fn rustup_lib() -> String {
    es::shell("rustc --print sysroot") + "/lib"
}

// this will be initially written to ~/.cargo/.runner/prelude and
// can then be edited.
const PRELUDE: &'static str = "
#![allow(unused_imports)]
#![allow(dead_code)]
use std::fs;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::env;
use std::path::{PathBuf,Path};
#[allow(unused_macros)]
macro_rules! debug {
    ($x:expr) => {
        println!(\"{} = {:?}\",stringify!($x),$x);
    }
}

";

fn runner_directory() -> PathBuf {
    let mut home = crate_utils::cargo_home();
    home.push(".runner");
    home
}

fn cargo(args: &[&str]) -> bool {
    let res = process::Command::new("cargo")
        .args(args)
        .status()
        .or_die("can't run cargo");
    res.success()
}

const STATIC_CACHE: &str = "static-cache";
const DYNAMIC_CACHE: &str = "dy-cache";

fn static_cache_dir() -> PathBuf {
    runner_directory().join(STATIC_CACHE)
}

fn static_cache_dir_check(args: &lapp::Args) -> PathBuf {
    let static_cache = static_cache_dir();
    if ! static_cache.exists() {
        args.quit("please build static cache with --create first");
    }
    static_cache
}

fn build_static_cache() -> bool {
    cargo(&["build"]) &&
    cargo(&["build","--release"]) &&
    cargo(&["doc"])
}

fn create_static_cache(crates: &[String], create: bool) {
    use std::io::prelude::*;
    let mut home = runner_directory();
    env::set_current_dir(&home).or_die("cannot change to home directory");
    if create {
        if ! cargo(&["new","--bin",STATIC_CACHE]) {
            es::quit("cannot create static cache");
        }
    }
    home.push(STATIC_CACHE);
    env::set_current_dir(&home).or_die("could not change to static cache directory");
    let tmpfile = env::temp_dir().join("Cargo.toml");
    fs::copy("Cargo.toml",&tmpfile).or_die("cannot back up Cargo.toml");
    {
        let mut deps = fs::OpenOptions::new().append(true)
            .open("Cargo.toml").or_die("could not append to Cargo.toml");
        for c in crates {
            write!(deps,"{}=\"*\"\n",c).or_die("could not modify Cargo.toml");
        }
    }
    if ! build_static_cache() {
        println!("Error occurred - restoring Cargo.toml");
        fs::copy(&tmpfile,"Cargo.toml").or_die("cannot restore Cargo.toml");
    }
}

// this is always called first and has the important role to ensure that
// runner's directory structure is created properly.
fn get_prelude() -> String {
    let home = runner_directory();
    let pristine = ! home.is_dir();
    if pristine {
        fs::create_dir(&home).or_die("cannot create runner directory");
    }
    let prelude = home.join("prelude");
    let bin = home.join("bin");
    if pristine {
        es::write_all(&prelude,PRELUDE);
        fs::create_dir(&home.join(DYNAMIC_CACHE)).or_die("cannot create dynamic cache");
    }
    if pristine || ! bin.is_dir() {
        fs::create_dir(&bin).or_die("cannot create output directory");
    }
    es::read_to_string(&prelude)
}

fn get_cache(build_static: bool, optimize: bool) -> PathBuf {
    let mut home = runner_directory();
    let scache;
    let cache = if build_static {
        let dir = if optimize {"release"} else {"debug"};
        scache = format!("{}/target/{}/deps",STATIC_CACHE,dir);
        &scache
    } else {DYNAMIC_CACHE};
    home.push(cache);
    home
}

fn add_aliases(aliases: Vec<String>) {
    if aliases.len() == 0 { return; }
    let alias_file = runner_directory().join("alias");
    let mut f = if alias_file.is_file() {
        fs::OpenOptions::new().append(true).open(&alias_file)
    } else {
        fs::File::create(&alias_file)
    }.or_die("cannot open runner alias file");

    for crate_alias in aliases {
        write!(f,"{}\n",crate_alias).or_die("cannot write to runner alias file");
    }
}

fn get_aliases() -> HashMap<String,String> {
    let alias_file = runner_directory().join("alias");
    if ! alias_file.is_file() { return HashMap::new(); }
	es::lines(es::open(&alias_file))
      .filter_map(|s| s.split_at_delim('=').trim()) // split into (String,String)
	  .to_map()
}

fn massage_snippet(code: String, prelude: String, extern_crates: Vec<String>, wild_crates: Vec<String>) -> String {
    fn indent_line(line: &str) -> String {
        format!("    {}\n",line)
    }
    let mut prefix = prelude;
    let mut body = String::new();
    {
        if extern_crates.len() > 0 {
            let aliases = get_aliases();
            for c in extern_crates {
                prefix += &if let Some(aliased) = aliases.get(&c) {
                    format!("extern crate {} as {};",aliased,c)
                } else {
                    format!("extern crate {};",c)
                };
            }
            for c in wild_crates {
                prefix += &format!("use {}::*;",c);
            }
        }
        let mut lines = code.lines();
        let mut first = true;
        for line in lines.by_ref() {
            let line = line.trim_left();
            if first { // files may start with #! shebang...
                if line.starts_with("#!") {
                    continue;
                }
                first = false;
            }
            if line.starts_with("//") || line.starts_with("#[macro_use") ||
                line.starts_with("extern ") || line.starts_with("use ") {
                prefix += line;
                prefix.push('\n');
            } else
            if line.len() > 0 {
                body += &indent_line(line);
                break;
            }
        }
        // and indent the rest!
        body.extend(lines.map(indent_line));
    }

    format!("{}
use std::error::Error;
fn run() -> Result<(),Box<Error>> {{
{}    Ok(())
}}
fn main() {{
    run().unwrap();
}}
        ",prefix,body)

}

const USAGE: &str = "
Compile and run small Rust snippets
  -s, --static build statically (default is dynamic)
  -O, --optimize optimized static build
  -e, --expression evaluate an expression
  -i, --iterator evaluate an iterator
  -n, --lines evaluate expression over stdin; 'line' is defined
  -x, --extern... (string) add an extern crate to the snippet
  -X, --wild... (string) like -x but implies wildcard use

  Cache Management:
  --create (string...) initialize the static cache with crates
  --add  (string...) add new crates to the cache (after --create)
  --edit  edit the static cache Cargo.toml
  --build rebuild the static cache
  --doc  display documentation
  --edit-prelude edit the default prelude for snippets
  --alias (string...) crate aliases in form alias=crate_name (used with -x)

  Dynamic compilation:
  -P, --crate-path show path of crate source in Cargo cache
  -C, --compile  compile crate dynamically (limited)
  --cfg... (string) pass configuration variables to rustc
  --libc  link dynamically against libc (special case)

  <program> (string) Rust program, snippet or expression
  <args> (string...) arguments to pass to program
";

fn main() {
    let args = lapp::parse_args(USAGE);
    let prelude = get_prelude();

    let aliases = args.get_strings("alias");
    if aliases.len() > 0 {
        add_aliases(aliases);
        return;
    }

    if args.get_bool("edit-prelude") {
        let rdir = runner_directory().join("prelude");
        edit(&rdir);
        return;
    }

    // Static Cache Management
    let crates = args.get_strings("create");
    if crates.len() > 0 {
        create_static_cache(&crates,true);
        return;
    }
    let crates = args.get_strings("add");
    if crates.len() > 0 {
        static_cache_dir_check(&args);
        create_static_cache(&crates,false);
        return;
    }

    if args.get_bool("edit") || args.get_bool("build") || args.get_bool("doc") {
        let static_cache = static_cache_dir_check(&args);
        if args.get_bool("build") {
            env::set_current_dir(&static_cache).or_die("static cache wasn't a directory?");
            build_static_cache();
        } else
        if args.get_bool("doc") {
            let docs = static_cache.join("target/doc/static_cache/index.html");
            open(&docs);
        } else {
            let toml = static_cache.join("Cargo.toml");
            edit(&toml);
        }
        return;
    }

    let first_arg = args.get_string("program");
    let file = PathBuf::from(&first_arg);

    // Dynamically linking crates (experimental!)
    if args.get_bool("crate-path") || args.get_bool("compile") {
        let (crate_path,crate_name) = if file.exists() {
            let filename = crate_utils::path_file_name(&file);
            if file.is_dir() { // assumed to be Cargo directory
                if ! file.join("Cargo.toml").exists() {
                    args.quit(&format!("not a Cargo project directory: {}",file.display()));
                }
                (file.join("src/lib.rs"), filename)
            } else { // should be just a Rust source file
                if file.extension().or_die("expecting extension") != "rs" {
                    args.quit("expecting Rust source file");
                }
                let name = crate_utils::path_file_name(&file.with_extension(""));
                (file, name)
            }
        } else {
            (crate_utils::cache_path(&first_arg).join("src/lib.rs"), first_arg)
        };
        if args.get_bool("crate-path") {
            println!("{}",crate_utils::cache_path(&crate_name).display());
        } else {
            let valid_crate_name = crate_name.replace('-',"_");
            let cache = get_cache(false, false);
            let mut builder = process::Command::new("rustc");
                builder.args(&["-C","prefer-dynamic"]).args(&["-C","debuginfo=0"])
                .arg("-L").arg(&cache)
                .args(&["--crate-type","dylib"])
                .arg("--out-dir").arg(&cache)
                .arg("--crate-name").arg(&valid_crate_name)
                .arg(crate_path);
           for c in args.get_strings("cfg") {
                builder.arg("--cfg").arg(&c);
           }
           if args.get_bool("libc") {
                builder.arg("--extern").arg(&format!("libc={}/liblibc.{}",cache.display(),SO));
           }
           builder.status().or_die("can't run rustc");
        }
        return;
    }

    let build_static = args.get_bool("static");
    let optimize = args.get_bool("optimize");

    // we'll pass rest of arguments to program
    let program_args = args.get_strings("args");

    let cache = get_cache(build_static, optimize);

    let mut snippet = false;
    let mut code = if args.get_bool("expression") {
        // Evaluating an expression: just print it out.
        format!("println!(\"{{:?}}\",{});", first_arg)
    } else
    if args.get_bool("iterator") {
        // The expression is anything that implements IntoIterator
        format!("let iter = {};\n for val in iter {{ println!(\"{{:?}}\",val);}}", first_arg)
    } else
    if args.get_bool("lines") {
        // The variable 'line' is available to an expression, evaluated for each line in stdin
        format!("
            let stdin = io::stdin();
            for line in stdin.lock().lines() {{
                let line = line?;
                let val = {};
                println!(\"{{:?}}\",val);
            }}
            ", first_arg)
    } else { // otherwise, just a file
        snippet = true;
        es::read_to_string(&file)
    };

    // proper Rust programs are accepted (this is a bit rough)
    if code.find("fn main").is_none() {
        let mut extern_crates = args.get_strings("extern");
        let wild_crates = args.get_strings("wild");
        if wild_crates.len() > 0 {
            extern_crates.extend(wild_crates.iter().cloned());
        }
        code = massage_snippet(code,prelude, extern_crates, wild_crates);
    }

    // we are going to put the expanded source and resulting exe in the runner bin dir
    let mut out_file = runner_directory().join("bin");
    if snippet {
        out_file.push(&file);
        out_file.set_extension("rs");
    } else {
        out_file.push("tmp.rs");
    }
    let mut program = out_file.clone();
    program.set_extension(EXE);

    es::write_all(&out_file,&code);

    let mut builder = process::Command::new("rustc");
    if ! build_static {
        builder.args(&["-C","prefer-dynamic"]).args(&["-C","debuginfo=0"]);
    } else {
        builder.arg(if optimize {"-O"} else {"-g"});
    }
    builder.arg("-L").arg(&cache);
    let status = builder.arg("-o").arg(&program)
        .arg(&out_file)
        .status().or_die("can't run rustc");
    if ! status.success() {
        return;
    }

    let mut builder = process::Command::new(&program);
    if ! build_static {
        builder.env("LD_LIBRARY_PATH",format!("{}:{}",rustup_lib(),cache.display()));
    }
    builder.args(&program_args)
        .status()
        .or_die(&format!("can't run program {:?}",program));

}

<|endoftext|>


<commit_before>use std::io::{self, BufReader, Write};
use std::io::BufRead;
use std::net::TcpStream;
use std::fs::File;

extern crate chrono;
use chrono::*;

extern crate stats;
use stats::median;

extern crate clap;
use clap::{Arg, App, ArgMatches};

mod percentile;
use percentile::percentile;

mod http_status;

mod request_response;
use request_response::*;

pub fn parse_logfile(path: &str, time_filter: Option<Duration>) -> Result<(Vec<Request>,Vec<Response>), io::Error> {
    let f = try!(File::open(path));

    let f = BufReader::new(f);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();

    for line in f.lines() {
        let line_value = &line.unwrap();

        if line_value.contains("->") {
            let r = try!(Request::new_from_log_line(&line_value));

            if time_filter.is_none() ||
              (time_filter.is_some() && r.is_between_times(UTC::now().with_timezone(&r.time.timezone()) - time_filter.unwrap(), UTC::now().with_timezone(&r.time.timezone()))) {
                requests.push(r);
            }
        }

        if line_value.contains("<-") {
            let r = try!(Response::new_from_log_line(&line_value));
            responses.push(r);
        }

    }

    responses.sort_by_key(|r| r.id);

    Ok((requests, responses))
}

#[derive(Eq, PartialEq)]
#[derive(Debug)]
pub struct RequestLogAnalyzerResult {
    count: usize,
    max: usize,
    min: usize,
    avg: usize,
    median: usize,
    percentile90: usize,
}

pub fn analyze(request_response_pairs: &Vec<RequestResponsePair>) -> Option<RequestLogAnalyzerResult> {
    if request_response_pairs.len() == 0 {
        return None;
    }

    let times: Vec<i64> = request_response_pairs.iter()
        .map(|rr: &RequestResponsePair| -> i64 {rr.response.response_time.num_milliseconds() })
        .collect();

    let sum: usize = times.iter().sum::<i64>() as usize;
    let avg: usize = sum / times.len();

    let max: usize = *times.iter().max().unwrap() as usize;
    let min: usize = *times.iter().min().unwrap() as usize;

    let percentile90: usize = percentile(&times, 0.9) as usize;

    let median = median(times.into_iter()).unwrap() as usize;

    Some(RequestLogAnalyzerResult {
        count: request_response_pairs.len().into(),
        max: max,
        min: min,
        avg: avg,
        median: median,
        percentile90: percentile90,
    })
}

fn render_terminal(result: RequestLogAnalyzerResult) {
    println!("count:\t{}", result.count);
    println!("time.avg:\t{}", result.avg);
    println!("time.min:\t{}", result.min);
    println!("time.median:\t{}", result.median);
    println!("time.90percent:\t{}", result.percentile90);
    println!("time.max:\t{}", result.max);
}

pub fn render_graphite<T: Write>(result: RequestLogAnalyzerResult, time: DateTime<FixedOffset>, prefix: Option<&str>, mut stream: T) {
    let prefix_text: &str;
    let prefix_separator: &str;

    match prefix {
        Some(p) => {
            prefix_text = p;
            prefix_separator = ".";
        }
        None => {
            prefix_text = "";
            prefix_separator = "";
        }
    };

    let mut write = |text: String| {
        stream.write(
            format!("{}{}{} {}\n", prefix_text, prefix_separator, text, time.timestamp() )
            .as_bytes()
        );
    };

    write(format!("requests.count {}", result.count));
    write(format!("requests.time.max {}", result.max));
    write(format!("requests.time.min {}", result.min));
    write(format!("requests.time.avg {}", result.avg));
    write(format!("requests.time.median {}", result.median));
    write(format!("requests.time.90percent {}", result.percentile90));
}

fn parse_args<'a>() -> ArgMatches<'a> {
    App::new("Request.log Analyzer")
        .arg(Arg::with_name("filename")
            .index(1)
            .value_name("FILE")
            .required(true)
            .help("Log file to analyze")
            .takes_value(true))
        .arg(Arg::with_name("time_filter_minutes")
            .value_name("MINUTES")
            .short("t")
            .help("Limit to the last n minutes")
            .takes_value(true))
        .arg(Arg::with_name("include_term")
            .value_name("TERM")
            .long("include")
            .help("Only includes lines that contain this term")
            .takes_value(true))
        .arg(Arg::with_name("exclude_term")
            .value_name("TERM")
            .long("exclude")
            .help("Excludes lines that contain this term")
            .takes_value(true))
        .arg(Arg::with_name("graphite-server")
            .value_name("GRAPHITE_SERVER")
            .long("graphite-server")
            .help("Send values to this Graphite server instead of stdout")
            .takes_value(true))
        .arg(Arg::with_name("graphite-port")
            .value_name("GRAPHITE_PORT")
            .long("graphite-port")
            .takes_value(true)
            .default_value("2003"))
        .arg(Arg::with_name("graphite-prefix")
            .value_name("GRAPHITE_PREFIX")
            .long("graphite-prefix")
            .help("Prefix for Graphite key, e.g. 'servers.prod.publisher1'")
            .takes_value(true))
        .get_matches()
}

fn main() {
    let args = parse_args();

    let filename = args.value_of("filename").unwrap();

    let time_filter = match args.value_of("time_filter_minutes") {
        Some(minutes) => Some(Duration::minutes(minutes.parse().unwrap())),
        None => None
    };

    let lines = parse_logfile(filename, time_filter);
    let (requests, responses) = lines.unwrap();

    let time_zone = &requests[0].time.timezone();

    let pairs: Vec<RequestResponsePair> = pair_requests_responses(requests, responses)
        .into_iter()
        .filter(|rr| rr.matches_include_exclude_filter(
            args.value_of("include_term"),
            args.value_of("exclude_term")
        ))
        .collect();

    if args.is_present("graphite-server") {
        let stream = TcpStream::connect(
            (
                args.value_of("graphite-server").unwrap(),
                args.value_of("graphite-port").unwrap().parse().unwrap()
            )
        ).unwrap();

        match analyze(&pairs) {
            Some(result) => render_graphite(result, UTC::now().with_timezone(time_zone), args.value_of("graphite-prefix"), stream),
            None => println!("No matching log lines in file.")
        }
    } else {
        match analyze(&pairs) {
            Some(result) => render_terminal(result),
            None => println!("No matching log lines in file.")
        }
    }
}

#[cfg(test)]
mod tests {
	use super::*;
    use request_response::*;
    extern crate chrono;
    use chrono::*;
    use std::str;
    use std::io::prelude::Write;
    use std::io::{self, BufReader};

    #[test]
    fn test_parse_logfile() {
        let lines = parse_logfile("src/test/simple-1.log", None);
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 2);
        assert_eq!(responses.len(), 2);
    }

    #[test]
    fn test_open_logfile_time_filter() {
        let time_filter: Duration = Duration::minutes(1);
        let lines = parse_logfile("src/test/simple-1.log", Some(time_filter));
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 0);

        let time_filter: Duration = Duration::minutes(52560000); // 100 years
        let lines = parse_logfile("src/test/simple-1.log", Some(time_filter));
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 2);
    }

    #[test]
    fn test_pair_requests_resonses() {
        let lines = parse_logfile("src/test/simple-1.log", None);
        let (requests, responses) = lines.unwrap();

        let result = pair_requests_responses(requests, responses);

        assert_eq!(result.len(), 2);

        assert_eq!(result[0].request.id, result[0].response.id);
        assert_eq!(result[1].request.id, result[1].response.id);
    }

    #[test]
    fn test_request_log_analyzer_result() {
        let lines = parse_logfile("src/test/response-time-calculations.log", None);
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result = analyze(&request_response_pairs);

        let expected = Some(RequestLogAnalyzerResult {
            count: 3,
            max: 100,
            min: 1,
            avg: 37,
            median: 10,
            percentile90: 100,
        });

        assert_eq!(result, expected);
    }

    #[test]
    fn test_request_log_analyze_none_matching() {
        let lines = parse_logfile("src/test/simple-1.log", Some(Duration::minutes(0)));
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result = analyze(&request_response_pairs);

        let expected = None;

        assert_eq!(result, expected);
    }

    #[test]
    fn test_90_percentile_calculation() {
        let lines = parse_logfile("src/test/percentile.log", None);
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result: RequestLogAnalyzerResult = analyze(&request_response_pairs).unwrap();

        assert_eq!(result.percentile90, 9);
    }

    struct MockTcpStream {
        write_calls: Vec<String>,
    }

    impl Write for MockTcpStream {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.write_calls.push(str::from_utf8(buf).unwrap().to_string());
            Ok(1)
        }

        fn flush(&mut self) -> io::Result<()> { Ok(()) }
    }

    #[test]
    fn test_render_graphite() {
        let mut mock_tcp_stream = MockTcpStream {
            write_calls: vec![]
        };

        render_graphite(RequestLogAnalyzerResult {
                count: 3,
                max: 100,
                min: 1,
                avg: 37,
                median: 10,
                percentile90: 100,
            },
            DateTime::parse_from_str("22/Sep/2016:22:41:59 +0200", "%d/%b/%Y:%H:%M:%S %z").unwrap(),
            None,
            &mut mock_tcp_stream
        );

        assert_eq!(&mock_tcp_stream.write_calls[0], "requests.count 3 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[1], "requests.time.max 100 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[2], "requests.time.min 1 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[3], "requests.time.avg 37 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[4], "requests.time.median 10 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[5], "requests.time.90percent 100 1474576919\n");
    }

    #[test]
    fn test_render_graphite_prefix() {
        let mut mock_tcp_stream = MockTcpStream {
            write_calls: vec![]
        };

        render_graphite(RequestLogAnalyzerResult {
                count: 3,
                max: 100,
                min: 1,
                avg: 37,
                median: 10,
                percentile90: 100,
            },
            DateTime::parse_from_str("22/Sep/2016:22:41:59 +0200", "%d/%b/%Y:%H:%M:%S %z").unwrap(),
            Some("my.prefix"),
            &mut mock_tcp_stream
        );

        assert_eq!(&mock_tcp_stream.write_calls[0], "my.prefix.requests.count 3 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[1], "my.prefix.requests.time.max 100 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[2], "my.prefix.requests.time.min 1 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[3], "my.prefix.requests.time.avg 37 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[4], "my.prefix.requests.time.median 10 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[5], "my.prefix.requests.time.90percent 100 1474576919\n");
    }
}
<commit_msg>Extract file opening with error handling<commit_after>use std::io::{self, BufReader, Write};
use std::io::BufRead;
use std::net::TcpStream;
use std::fs::File;

extern crate chrono;
use chrono::*;

extern crate stats;
use stats::median;

extern crate clap;
use clap::{Arg, App, ArgMatches};

mod percentile;
use percentile::percentile;

mod http_status;

mod request_response;
use request_response::*;

fn open_logfile(path: &str) -> BufReader<File> {
    let file = File::open(path);

    match file {
        Ok(f) => BufReader::new(f),
        Err(err) => panic!("Could not open file {}: {}", path, err)
    }
}

pub fn parse_logfile(path: &str, time_filter: Option<Duration>) -> Result<(Vec<Request>,Vec<Response>), io::Error> {
    let f = open_logfile(path);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();

    for line in f.lines() {
        let line_value = &line.unwrap();

        if line_value.contains("->") {
            let r = try!(Request::new_from_log_line(&line_value));

            if time_filter.is_none() ||
              (time_filter.is_some() && r.is_between_times(UTC::now().with_timezone(&r.time.timezone()) - time_filter.unwrap(), UTC::now().with_timezone(&r.time.timezone()))) {
                requests.push(r);
            }
        }

        if line_value.contains("<-") {
            let r = try!(Response::new_from_log_line(&line_value));
            responses.push(r);
        }

    }

    responses.sort_by_key(|r| r.id);

    Ok((requests, responses))
}

#[derive(Eq, PartialEq)]
#[derive(Debug)]
pub struct RequestLogAnalyzerResult {
    count: usize,
    max: usize,
    min: usize,
    avg: usize,
    median: usize,
    percentile90: usize,
}

pub fn analyze(request_response_pairs: &Vec<RequestResponsePair>) -> Option<RequestLogAnalyzerResult> {
    if request_response_pairs.len() == 0 {
        return None;
    }

    let times: Vec<i64> = request_response_pairs.iter()
        .map(|rr: &RequestResponsePair| -> i64 {rr.response.response_time.num_milliseconds() })
        .collect();

    let sum: usize = times.iter().sum::<i64>() as usize;
    let avg: usize = sum / times.len();

    let max: usize = *times.iter().max().unwrap() as usize;
    let min: usize = *times.iter().min().unwrap() as usize;

    let percentile90: usize = percentile(&times, 0.9) as usize;

    let median = median(times.into_iter()).unwrap() as usize;

    Some(RequestLogAnalyzerResult {
        count: request_response_pairs.len().into(),
        max: max,
        min: min,
        avg: avg,
        median: median,
        percentile90: percentile90,
    })
}

fn render_terminal(result: RequestLogAnalyzerResult) {
    println!("count:\t{}", result.count);
    println!("time.avg:\t{}", result.avg);
    println!("time.min:\t{}", result.min);
    println!("time.median:\t{}", result.median);
    println!("time.90percent:\t{}", result.percentile90);
    println!("time.max:\t{}", result.max);
}

pub fn render_graphite<T: Write>(result: RequestLogAnalyzerResult, time: DateTime<FixedOffset>, prefix: Option<&str>, mut stream: T) {
    let prefix_text: &str;
    let prefix_separator: &str;

    match prefix {
        Some(p) => {
            prefix_text = p;
            prefix_separator = ".";
        }
        None => {
            prefix_text = "";
            prefix_separator = "";
        }
    };

    let mut write = |text: String| {
        stream.write(
            format!("{}{}{} {}\n", prefix_text, prefix_separator, text, time.timestamp() )
            .as_bytes()
        );
    };

    write(format!("requests.count {}", result.count));
    write(format!("requests.time.max {}", result.max));
    write(format!("requests.time.min {}", result.min));
    write(format!("requests.time.avg {}", result.avg));
    write(format!("requests.time.median {}", result.median));
    write(format!("requests.time.90percent {}", result.percentile90));
}

fn parse_args<'a>() -> ArgMatches<'a> {
    App::new("Request.log Analyzer")
        .arg(Arg::with_name("filename")
            .index(1)
            .value_name("FILE")
            .required(true)
            .help("Log file to analyze")
            .takes_value(true))
        .arg(Arg::with_name("time_filter_minutes")
            .value_name("MINUTES")
            .short("t")
            .help("Limit to the last n minutes")
            .takes_value(true))
        .arg(Arg::with_name("include_term")
            .value_name("TERM")
            .long("include")
            .help("Only includes lines that contain this term")
            .takes_value(true))
        .arg(Arg::with_name("exclude_term")
            .value_name("TERM")
            .long("exclude")
            .help("Excludes lines that contain this term")
            .takes_value(true))
        .arg(Arg::with_name("graphite-server")
            .value_name("GRAPHITE_SERVER")
            .long("graphite-server")
            .help("Send values to this Graphite server instead of stdout")
            .takes_value(true))
        .arg(Arg::with_name("graphite-port")
            .value_name("GRAPHITE_PORT")
            .long("graphite-port")
            .takes_value(true)
            .default_value("2003"))
        .arg(Arg::with_name("graphite-prefix")
            .value_name("GRAPHITE_PREFIX")
            .long("graphite-prefix")
            .help("Prefix for Graphite key, e.g. 'servers.prod.publisher1'")
            .takes_value(true))
        .get_matches()
}

fn main() {
    let args = parse_args();

    let filename = args.value_of("filename").unwrap();

    let time_filter = match args.value_of("time_filter_minutes") {
        Some(minutes) => Some(Duration::minutes(minutes.parse().unwrap())),
        None => None
    };

    let lines = parse_logfile(filename, time_filter);
    let (requests, responses) = lines.unwrap();

    let time_zone = &requests[0].time.timezone();

    let pairs: Vec<RequestResponsePair> = pair_requests_responses(requests, responses)
        .into_iter()
        .filter(|rr| rr.matches_include_exclude_filter(
            args.value_of("include_term"),
            args.value_of("exclude_term")
        ))
        .collect();

    if args.is_present("graphite-server") {
        let stream = TcpStream::connect(
            (
                args.value_of("graphite-server").unwrap(),
                args.value_of("graphite-port").unwrap().parse().unwrap()
            )
        ).unwrap();

        match analyze(&pairs) {
            Some(result) => render_graphite(result, UTC::now().with_timezone(time_zone), args.value_of("graphite-prefix"), stream),
            None => println!("No matching log lines in file.")
        }
    } else {
        match analyze(&pairs) {
            Some(result) => render_terminal(result),
            None => println!("No matching log lines in file.")
        }
    }
}

#[cfg(test)]
mod tests {
	use super::*;
    use request_response::*;
    extern crate chrono;
    use chrono::*;
    use std::str;
    use std::io::prelude::Write;
    use std::io::{self, BufReader};

    #[test]
    fn test_parse_logfile() {
        let lines = parse_logfile("src/test/simple-1.log", None);
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 2);
        assert_eq!(responses.len(), 2);
    }

    #[test]
    fn test_open_logfile_time_filter() {
        let time_filter: Duration = Duration::minutes(1);
        let lines = parse_logfile("src/test/simple-1.log", Some(time_filter));
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 0);

        let time_filter: Duration = Duration::minutes(52560000); // 100 years
        let lines = parse_logfile("src/test/simple-1.log", Some(time_filter));
        let (requests, responses) = lines.unwrap();

        assert_eq!(requests.len(), 2);
    }

    #[test]
    fn test_pair_requests_resonses() {
        let lines = parse_logfile("src/test/simple-1.log", None);
        let (requests, responses) = lines.unwrap();

        let result = pair_requests_responses(requests, responses);

        assert_eq!(result.len(), 2);

        assert_eq!(result[0].request.id, result[0].response.id);
        assert_eq!(result[1].request.id, result[1].response.id);
    }

    #[test]
    fn test_request_log_analyzer_result() {
        let lines = parse_logfile("src/test/response-time-calculations.log", None);
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result = analyze(&request_response_pairs);

        let expected = Some(RequestLogAnalyzerResult {
            count: 3,
            max: 100,
            min: 1,
            avg: 37,
            median: 10,
            percentile90: 100,
        });

        assert_eq!(result, expected);
    }

    #[test]
    fn test_request_log_analyze_none_matching() {
        let lines = parse_logfile("src/test/simple-1.log", Some(Duration::minutes(0)));
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result = analyze(&request_response_pairs);

        let expected = None;

        assert_eq!(result, expected);
    }

    #[test]
    fn test_90_percentile_calculation() {
        let lines = parse_logfile("src/test/percentile.log", None);
        let (requests, responses) = lines.unwrap();

        let request_response_pairs = pair_requests_responses(requests, responses);

        let result: RequestLogAnalyzerResult = analyze(&request_response_pairs).unwrap();

        assert_eq!(result.percentile90, 9);
    }

    struct MockTcpStream {
        write_calls: Vec<String>,
    }

    impl Write for MockTcpStream {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.write_calls.push(str::from_utf8(buf).unwrap().to_string());
            Ok(1)
        }

        fn flush(&mut self) -> io::Result<()> { Ok(()) }
    }

    #[test]
    fn test_render_graphite() {
        let mut mock_tcp_stream = MockTcpStream {
            write_calls: vec![]
        };

        render_graphite(RequestLogAnalyzerResult {
                count: 3,
                max: 100,
                min: 1,
                avg: 37,
                median: 10,
                percentile90: 100,
            },
            DateTime::parse_from_str("22/Sep/2016:22:41:59 +0200", "%d/%b/%Y:%H:%M:%S %z").unwrap(),
            None,
            &mut mock_tcp_stream
        );

        assert_eq!(&mock_tcp_stream.write_calls[0], "requests.count 3 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[1], "requests.time.max 100 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[2], "requests.time.min 1 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[3], "requests.time.avg 37 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[4], "requests.time.median 10 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[5], "requests.time.90percent 100 1474576919\n");
    }

    #[test]
    fn test_render_graphite_prefix() {
        let mut mock_tcp_stream = MockTcpStream {
            write_calls: vec![]
        };

        render_graphite(RequestLogAnalyzerResult {
                count: 3,
                max: 100,
                min: 1,
                avg: 37,
                median: 10,
                percentile90: 100,
            },
            DateTime::parse_from_str("22/Sep/2016:22:41:59 +0200", "%d/%b/%Y:%H:%M:%S %z").unwrap(),
            Some("my.prefix"),
            &mut mock_tcp_stream
        );

        assert_eq!(&mock_tcp_stream.write_calls[0], "my.prefix.requests.count 3 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[1], "my.prefix.requests.time.max 100 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[2], "my.prefix.requests.time.min 1 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[3], "my.prefix.requests.time.avg 37 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[4], "my.prefix.requests.time.median 10 1474576919\n");
        assert_eq!(&mock_tcp_stream.write_calls[5], "my.prefix.requests.time.90percent 100 1474576919\n");
    }
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (eliza@elizas.website)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! # SOS kernel
//! This crate contains the kernel for SOS, the Stupid Operating System.
//!
//! # SOS: the Stupid Operating System
//! SOS is a simple, tiny toy OS implemented in Rust. It targets the `x86`,
//! `x86_64`, and ARM v7 CPU architectures.
//!
//! I'm writing this mostly for fun, to learn more about OS design and kernel
//! hacking, so don't expect anything new or exciting out of this project.
//!
//! SOS is copyright 2015-2017 Eliza Weisman, and is released under the terms
//! of the MIT license.

#![crate_name = "sos_kernel"]

#![doc(html_root_url = "https://hawkw.github.io/sos-kernel/")]

#![feature( lang_items, asm, naked_functions )]
#![feature( linkage )]
#![feature( const_fn
          , slice_patterns
          , associated_consts
          , type_ascription
          , custom_derive )]
#![feature(collections)]

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr( any(target_arch = "x86_64", target_arch="x86")
           , feature(abi_x86_interrupt))]

#![no_std]
#![cfg_attr(not(test), no_main)]

// -- non-SOS dependencies --------------------------------------------------
#[macro_use] extern crate lazy_static;
#[macro_use] extern crate bitflags;
#[macro_use] extern crate log;

extern crate collections;
extern crate rlibc;
extern crate spin;

// -- SOS dependencies ------------------------------------------------------
#[macro_use] extern crate vga;

extern crate alloc;
extern crate cpu;
extern crate elf;
extern crate paging;
extern crate params;
extern crate memory;
extern crate util;

#[macro_use] pub mod io;

pub mod heap;
pub mod arch;
pub mod logger;

use params::InitParams;

#[macro_use]
macro_rules! init_log {
    (f$dots:expr, $task:expr, $msg:expr) => (
        println!( "{task:<40}{res:>38}\n{msg:>.width$}"
                , task = format!("{:>.width$}", $task, width = $dots)
                , res = "[ FAIL ]"
                , msg = $msg
                , width = $dots + 1
                )
    );
    (fail: $dots:expr, $task:expr) => (
            println!( "{task:<40}{res:>38}"
                    , task = format!("{:>.width$}", $task, width = $dots)
                    , res = "[ FAIL ]"
                    )
    );
    (okay: $dots:expr, $task:expr, $msg:expr) => (
        println!( "{task:<40}{res:>38}\n{msg:>.width$}"
                , task = format!("{:>.width$}", $task, width = $dots)
                , res = "[ OKAY ]"
                , msg = $msg
                , width = $dots + 1
                )
    );
    (okay: $dots:expr, $task:expr) => (
            println!( "{task:<40}{res:>38}"
                    , task = format!("{:>.width$}", $task, width = $dots)
                    , res = "[ OKAY ]"
                    )
    );
}

macro_rules! attempt {
    ($task:expr => $msg:expr, dots: $dots:expr) => {
        {
            print!("{}{}", $dots, $msg);
            match $task {
               Ok(result) => {
                    println!( "{:indent$}[ OKAY ]"
                            , indent=80 - concat!($dots,$msg).len());
                    info!("{} [ OKAY ]", $msg);
                    result
                }
              , Err(why) => {
                    println!( "{:indent$}[ FAIL ]"
                            , indent=80 - concat!($dots,$msg).len());
                    panic!("{:?}", why);
              }
            }
        }

    };
}

/// SOS version number
pub const VERSION_STRING: &'static str
    = concat!("Stupid Operating System v", env!("CARGO_PKG_VERSION"));

/// Kernel main loop
pub fn kernel_main() -> ! {
    let mut a_vec = collections::vec::Vec::<usize>::new();
    info!(target: "test", "Created a vector in kernel space! {:?}", a_vec);
    a_vec.push(1);
    info!(target: "test", "pushed to vec: {:?}", a_vec);
    a_vec.push(2);
    info!(target: "test", "pushed to vec: {:?}", a_vec);

    // let mut frame_allocator = frame_alloc::FrameAllocator::new();
    // paging::test_paging(&mut frame_allocator);

    loop { }
}

/// Kernel initialization function called into by architecture-specific init
///
/// Our initialization process essentially looks like this:
///
/// ```text
/// +-------------+
/// | bootloader  |
/// | (multiboot) |
/// +------|------+
/// +------V------+
/// | start.asm   |
/// +------|------+
/// +------|--------------------------------------------------------+
/// |      |           RUST-LAND KERNEL FUNCTIONS                   |
/// |      V                                                        |
/// | arch_init() ----------> kernel_init() --------> kernel_main() |
/// | + collects boot info   + initializes interrupts               |
/// |   from arch-specific   + initializes the heap                 |
/// |   sources              + remaps the kernel into the higher    |
/// | + some CPU-specific      half of the address space            |
/// |   configuration                                               |
/// +---------------------------------------------------------------+
/// ```
pub fn kernel_init(params: &InitParams) {
    kinfoln!("Hello from the kernel!");
    kinfoln!("Got init params: {:#?}", params );
    // -- initialize interrupts ----------------------------------------------
    kinfoln!(dots: " . ", "Initializing interrupts:");
    // TODO: this whole function *may* want to just be made `unsafe`...
    unsafe { arch::interrupts::initialize(); };
    kinfoln!(dots: " . ", target: "Enabling interrupts", "[ OKAY ]");

    // -- initialize the heap ------------------------------------------------
    kinfoln!(dots: " . ", "Preparing to initialize heap.");
    if unsafe { heap::initialize(params) }.is_ok() {
        kinfoln!( dots: " . ", target: "Intializing heap"
                , "[ OKAY ]"
                );
        kinfoln!( dots: " . . "
                , "Heap begins at {:#x} and ends at {:#x}"
                , params.heap_base, params.heap_top);
    } else {
        kinfoln!( dots: " . ", target: "Intializing heap"
                , "[ FAIL ]"
                );
    }

    // -- remap the kernel ----------------------------------------------------
    kinfoln!(dots: " . ", "Remapping the kernel:");

    // let frame_allocator = alloc::buddy::BuddyFrameAllocator::new();

    // ::paging::kernel_remap(&params, &frame_allocator);

    kinfoln!( dots: " . ", target: "Remapping the kernel", "[ OKAY ]");

    println!("\n{} {}-bit\n", VERSION_STRING, arch::ARCH_BITS);

    // -- call into kernel main loop ------------------------------------------
    // (currently, this does nothing)
    kernel_main()
}

/// This fake `main` function exists only to placate `cargo test`.
#[cfg(test)]
fn main() {

}
<commit_msg>feat(kernel_init): nicer VGA logging for kernel remap<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (eliza@elizas.website)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! # SOS kernel
//! This crate contains the kernel for SOS, the Stupid Operating System.
//!
//! # SOS: the Stupid Operating System
//! SOS is a simple, tiny toy OS implemented in Rust. It targets the `x86`,
//! `x86_64`, and ARM v7 CPU architectures.
//!
//! I'm writing this mostly for fun, to learn more about OS design and kernel
//! hacking, so don't expect anything new or exciting out of this project.
//!
//! SOS is copyright 2015-2017 Eliza Weisman, and is released under the terms
//! of the MIT license.

#![crate_name = "sos_kernel"]

#![doc(html_root_url = "https://hawkw.github.io/sos-kernel/")]

#![feature( lang_items, asm, naked_functions )]
#![feature( linkage )]
#![feature( const_fn
          , slice_patterns
          , associated_consts
          , type_ascription
          , custom_derive )]
#![feature(collections)]

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr( any(target_arch = "x86_64", target_arch="x86")
           , feature(abi_x86_interrupt))]

#![no_std]
#![cfg_attr(not(test), no_main)]

// -- non-SOS dependencies --------------------------------------------------
#[macro_use] extern crate lazy_static;
#[macro_use] extern crate bitflags;
#[macro_use] extern crate log;

extern crate collections;
extern crate rlibc;
extern crate spin;

// -- SOS dependencies ------------------------------------------------------
#[macro_use] extern crate vga;

extern crate alloc;
extern crate cpu;
extern crate elf;
extern crate paging;
extern crate params;
extern crate memory;
extern crate util;

#[macro_use] pub mod io;

pub mod heap;
pub mod arch;
pub mod logger;

use params::InitParams;

#[macro_use]
macro_rules! init_log {
    (fail: $dots:expr, $task:expr, $msg:expr) => (
        kinfoln!( dots: $dots, "{task:<40}{res:>38}"
                , task = $task
                , res = "[ FAIL ]"
                )
    );
    (fail: $dots:expr, $task:expr) => (
            kinfoln!( "{task:<40}{res:>38}"
                    , task = $task
                    , res = "[ FAIL ]"
                    )
    );
    (okay: $dots:expr, $task:expr, $msg:expr) => (
        println!( "{task:<40}{res:>38}\n{msg:>.width$}"
                , task = format!("{:>.width$}", $task, width = $dots)
                , res = "[ OKAY ]"
                , msg = $msg
                , width = $dots + 1
                )
    );
    (okay: $dots:expr, $task:expr) => (
            kinfoln!( dots: $dots, "{task:<40}{res:>38}"
                    , task = $task
                    , res = "[ OKAY ]"
                    )
    );
}

macro_rules! attempt {
    ($task:expr => $msg:expr, dots: $dots:expr) => {
        {
            print!("{}{}", $dots, $msg);
            match $task {
               Ok(result) => {
                    println!( "{:indent$}[ OKAY ]"
                            , indent=80 - concat!($dots,$msg).len());
                    info!("{} [ OKAY ]", $msg);
                    result
                }
              , Err(why) => {
                    println!( "{:indent$}[ FAIL ]"
                            , indent=80 - concat!($dots,$msg).len());
                    panic!("{:?}", why);
              }
            }
        }

    };
}

/// SOS version number
pub const VERSION_STRING: &'static str
    = concat!("Stupid Operating System v", env!("CARGO_PKG_VERSION"));

/// Kernel main loop
pub fn kernel_main() -> ! {
    // let mut a_vec = collections::vec::Vec::<usize>::new();
    // info!(target: "test", "Created a vector in kernel space! {:?}", a_vec);
    // a_vec.push(1);
    // info!(target: "test", "pushed to vec: {:?}", a_vec);
    // a_vec.push(2);
    // info!(target: "test", "pushed to vec: {:?}", a_vec);

    // let mut frame_allocator = frame_alloc::FrameAllocator::new();
    // paging::test_paging(&mut frame_allocator);

    loop { }
}

/// Kernel initialization function called into by architecture-specific init
///
/// Our initialization process essentially looks like this:
///
/// ```text
/// +-------------+
/// | bootloader  |
/// | (multiboot) |
/// +------|------+
/// +------V------+
/// | start.asm   |
/// +------|------+
/// +------|--------------------------------------------------------+
/// |      |           RUST-LAND KERNEL FUNCTIONS                   |
/// |      V                                                        |
/// | arch_init() ----------> kernel_init() --------> kernel_main() |
/// | + collects boot info   + initializes interrupts               |
/// |   from arch-specific   + initializes the heap                 |
/// |   sources              + remaps the kernel into the higher    |
/// | + some CPU-specific      half of the address space            |
/// |   configuration                                               |
/// +---------------------------------------------------------------+
/// ```
pub fn kernel_init(params: &InitParams) {
    use alloc::frame::mem_map::MemMapAllocator;
    use ::paging::kernel_remap;

    kinfoln!("Hello from the kernel!");
    // kinfoln!("Got init params: {:#?}", params );

    // -- remap the kernel ----------------------------------------------------
    let mut frame_allocator = MemMapAllocator::from(params);
    kinfoln!(dots: " . ", "Remapping the kernel...");
    let page_table = match kernel_remap(&params, &mut frame_allocator) {
        Ok(p) => { init_log!(okay: " . ", "Remapping the kernel"); p}
      , Err(why) => {
            init_log!(fail: " . ", "Remapping the kernel");
            panic!( "Could not remap kernel: {}", why)
        }
    };

    paging::test_paging(&mut frame_allocator);

    // -- initialize the heap ------------------------------------------------
    attempt!( unsafe { heap::initialize(params) } =>
             "Intializing heap...", dots: " . ");
    kinfoln!( dots: " . . "
            , "Heap begins at {:#x} and ends at {:#x}"
            , params.heap_base, params.heap_top);

    // -- initialize interrupts ----------------------------------------------
    // attempt!( unsafe { arch::interrupts::initialize() } =>
    //           "Initializing interrupts...", dots: " . " );

    println!("\n{} {}-bit\n", VERSION_STRING, arch::ARCH_BITS);

    // -- call into kernel main loop ------------------------------------------
    // (currently, this does nothing)
    kernel_main()
}

/// This fake `main` function exists only to placate `cargo test`.
#[cfg(test)]
fn main() {

}
<|endoftext|>


<commit_before>#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate itertools;
extern crate walkdir;

use std::env;
use std::process::{Command, Output};
use clap::{App, SubCommand, AppSettings};
use itertools::Itertools;
use walkdir::{DirEntry, WalkDirIterator};

fn announce(banner: &str) {
    let mut line = String::new();
    for _ in 0..banner.len() {
        line.push('-');
    }
    println!("{}", line);
    println!("{}", banner);
    println!("{}", line);
}

fn print_ident(buf: Vec<u8>) {
    for line in String::from_utf8_lossy(&buf[..]).lines() {
        println!("        {}", line);
    }
}

fn report_output(output: Output) {
    if output.status.success() {
        print_ident(output.stdout);
    } else {
        print_ident(output.stderr);
    }
    // I am still not sure what is more idiomatic - the 'if' above or the 'match' below
    //
    // match output.status.success() {
    //     true => print_ident(output.stdout),
    //     false => print_ident(output.stderr),
    // }
    println!("");
}

const CARGO: &'static str = "cargo";
const MIN_DEPTH: usize = 1;
const MAX_DEPTH: usize = 1;

fn main() {

    let matches = App::new(CARGO)
                      .bin_name(CARGO)
                      .version(crate_version!())
                      .about("Run cargo command on multiple crates")
                      .setting(AppSettings::SubcommandRequired)
                      .setting(AppSettings::ArgRequiredElseHelp)
                      .subcommand(SubCommand::with_name("multi")
                                      .version(crate_version!())
                                      .setting(AppSettings::ArgRequiredElseHelp)
                                      .setting(AppSettings::TrailingVarArg)
                                      .arg_from_usage("<cmd>... 'cargo command to run'"))
                      .get_matches();

    let mut cargo_cmd = Command::new(CARGO);
    let mut banner = String::from("Executing ") + CARGO;

    if let Some(arg_cmd) = matches.subcommand_matches("multi")
                                  .and_then(|m| m.values_of("cmd")) {
        for arg in arg_cmd {
            cargo_cmd.arg(arg);
            banner = banner + " " + arg;
        }
    }

    let is_crate = |e: &DirEntry| e.path().join("Cargo.toml").exists();
    let display_path = |e: &DirEntry| println!("{}:", e.file_name().to_string_lossy());
    let execute = |e: DirEntry| cargo_cmd.current_dir(e.path()).output().ok();

    if let Ok(cwd) = env::current_dir() {
        announce(&banner);
        walkdir::WalkDir::new(cwd)
            .min_depth(MIN_DEPTH)
            .max_depth(MAX_DEPTH)
            .into_iter()
            .filter_entry(is_crate)
            .filter_map(|e| e.ok())
            .inspect(display_path)
            .filter_map(execute)
            .foreach(report_output);
    }
}
<commit_msg>Simplify banner a bit<commit_after>#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate itertools;
extern crate walkdir;

use std::env;
use std::process::{Command, Output};
use clap::{App, SubCommand, AppSettings};
use itertools::Itertools;
use walkdir::{DirEntry, WalkDirIterator};

fn announce(banner: &str) {
    let mut line = String::new();
    for _ in 0..banner.len() {
        line.push('-');
    }
    println!("{}\n{}\n{}", line, banner, line);
}

fn print_ident(buf: Vec<u8>) {
    for line in String::from_utf8_lossy(&buf[..]).lines() {
        println!("        {}", line);
    }
}

fn report_output(output: Output) {
    if output.status.success() {
        print_ident(output.stdout);
    } else {
        print_ident(output.stderr);
    }
    // I am still not sure what is more idiomatic - the 'if' above or the 'match' below
    //
    // match output.status.success() {
    //     true => print_ident(output.stdout),
    //     false => print_ident(output.stderr),
    // }
    println!("");
}

const CARGO: &'static str = "cargo";
const MIN_DEPTH: usize = 1;
const MAX_DEPTH: usize = 1;

fn main() {

    let matches = App::new(CARGO)
                      .bin_name(CARGO)
                      .version(crate_version!())
                      .about("Run cargo command on multiple crates")
                      .setting(AppSettings::SubcommandRequired)
                      .setting(AppSettings::ArgRequiredElseHelp)
                      .subcommand(SubCommand::with_name("multi")
                                      .version(crate_version!())
                                      .setting(AppSettings::ArgRequiredElseHelp)
                                      .setting(AppSettings::TrailingVarArg)
                                      .arg_from_usage("<cmd>... 'cargo command to run'"))
                      .get_matches();

    let mut cargo_cmd = Command::new(CARGO);
    let mut banner = String::from("Executing ") + CARGO;

    if let Some(arg_cmd) = matches.subcommand_matches("multi")
                                  .and_then(|m| m.values_of("cmd")) {
        for arg in arg_cmd {
            cargo_cmd.arg(arg);
            banner = banner + " " + arg;
        }
    }

    let is_crate = |e: &DirEntry| e.path().join("Cargo.toml").exists();
    let display_path = |e: &DirEntry| println!("{}:", e.file_name().to_string_lossy());
    let execute = |e: DirEntry| cargo_cmd.current_dir(e.path()).output().ok();

    if let Ok(cwd) = env::current_dir() {
        announce(&banner);
        walkdir::WalkDir::new(cwd)
            .min_depth(MIN_DEPTH)
            .max_depth(MAX_DEPTH)
            .into_iter()
            .filter_entry(is_crate)
            .filter_map(|e| e.ok())
            .inspect(display_path)
            .filter_map(execute)
            .foreach(report_output);
    }
}
<|endoftext|>