<commit_before>//! The internal representation of a book and infrastructure for loading it from
//! disk and building it.
//!
//! For examples on using `MDBook`, consult the [top-level documentation][1].
//!
//! [1]: ../index.html

#![deny(missing_docs)]

mod summary;
mod book;
mod init;

pub use self::book::{load_book, Book, BookItem, BookItems, Chapter};
pub use self::summary::{parse_summary, Link, SectionNumber, Summary, SummaryItem};
pub use self::init::BookBuilder;

use std::path::PathBuf;
use std::io::Write;
use std::process::Command;
use tempdir::TempDir;
use toml::Value;

use utils;
use renderer::{CmdRenderer, HtmlHandlebars, RenderContext, Renderer};
use preprocess::{Preprocessor, LinkPreprocessor, PreprocessorContext};
use errors::*;

use config::Config;

/// The object used to manage and build a book.
pub struct MDBook {
    /// The book's root directory.
    pub root: PathBuf,
    /// The configuration used to tweak now a book is built.
    pub config: Config,
    /// A representation of the book's contents in memory.
    pub book: Book,
    renderers: Vec<Box<Renderer>>,

    /// The URL used for live reloading when serving up the book.
    pub livereload: Option<String>,

    /// List of pre-processors to be run on the book
    preprocessors: Vec<Box<Preprocessor>>
}

impl MDBook {
    /// Load a book from its root directory on disk.
    pub fn load<P: Into<PathBuf>>(book_root: P) -> Result<MDBook> {
        let book_root = book_root.into();
        let config_location = book_root.join("book.toml");

        // the book.json file is no longer used, so we should emit a warning to
        // let people know to migrate to book.toml
        if book_root.join("book.json").exists() {
            warn!("It appears you are still using book.json for configuration.");
            warn!("This format is no longer used, so you should migrate to the");
            warn!("book.toml format.");
            warn!("Check the user guide for migration information:");
            warn!("\thttps://rust-lang-nursery.github.io/mdBook/format/config.html");
        }

        let mut config = if config_location.exists() {
            debug!("[*] Loading config from {}", config_location.display());
            Config::from_disk(&config_location)?
        } else {
            Config::default()
        };

        config.update_from_env();

        if log_enabled!(::log::Level::Trace) {
            for line in format!("Config: {:#?}", config).lines() {
                trace!("{}", line);
            }
        }

        MDBook::load_with_config(book_root, config)
    }

    /// Load a book from its root directory using a custom config.
    pub fn load_with_config<P: Into<PathBuf>>(book_root: P, config: Config) -> Result<MDBook> {
        let root = book_root.into();

        let src_dir = root.join(&config.book.src);
        let book = book::load_book(&src_dir, &config.build)?;
        let livereload = None;

        let renderers = determine_renderers(&config);
        let preprocessors = determine_preprocessors(&config)?;

        Ok(MDBook {
            root,
            config,
            book,
            renderers,
            livereload,
            preprocessors,
        })
    }

    /// 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::book::BookItem;
    /// # #[allow(unused_variables)]
    /// # fn main() {
    /// # let book = MDBook::load("mybook").unwrap();
    /// for item in book.iter() {
    ///     match *item {
    ///         BookItem::Chapter(ref chapter) => {},
    ///         BookItem::Separator => {},
    ///     }
    /// }
    ///
    /// // 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 {
        self.book.iter()
    }

    /// `init()` gives you a `BookBuilder` which you can use to setup a new book
    /// and its accompanying directory structure.
    ///
    /// The `BookBuilder` 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 path provided as the root directory for your book, then adds
    /// in a `src/` directory containing a `SUMMARY.md` and `chapter_1.md` file
    /// to get you started.
    pub fn init<P: Into<PathBuf>>(book_root: P) -> BookBuilder {
        BookBuilder::new(book_root)
    }

    /// Tells the renderer to build our book and put it in the build directory.
    pub fn build(&self) -> Result<()> {
        debug!("[fn]: build");

        let mut preprocessed_book = self.book.clone();
        let preprocess_ctx = PreprocessorContext {
            src_dir: self.source_dir(),
        };

        for preprocessor in &self.preprocessors {
            preprocessor.run(&preprocess_ctx, &mut preprocessed_book)?;
        }

        for renderer in &self.renderers {
            self.run_renderer(&preprocessed_book, renderer.as_ref())?;
        }

        Ok(())
    }

    fn run_renderer(&self, preprocessed_book: &Book, renderer: &Renderer) -> Result<()> {
        let name = renderer.name();
        let build_dir = self.build_dir_for(name);
        if build_dir.exists() {
            debug!(
                "Cleaning build dir for the \"{}\" renderer ({})",
                name,
                build_dir.display()
            );

            utils::fs::remove_dir_content(&build_dir)
                .chain_err(|| "Unable to clear output directory")?;
        }

        let render_context = RenderContext::new(
            self.root.clone(),
            preprocessed_book.clone(),
            self.config.clone(),
            build_dir,
        );

        renderer
            .render(&render_context)
            .chain_err(|| "Rendering failed")
    }

    /// 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)
    pub fn with_renderer<R: Renderer + 'static>(&mut self, renderer: R) -> &mut Self {
        self.renderers.push(Box::new(renderer));
        self
    }

    /// You can add a new preprocessor by using this method.
    /// The only requirement is for your renderer to implement the Preprocessor trait.
    pub fn with_preprecessor<P: Preprocessor + 'static>(&mut self, preprocessor: P) -> &mut Self {
        self.preprocessors.push(Box::new(preprocessor));
        self
    }

    /// Run `rustdoc` tests on the book, linking against the provided libraries.
    pub fn test(&mut self, library_paths: Vec<&str>) -> Result<()> {
        let library_args: Vec<&str> = (0..library_paths.len())
            .map(|_| "-L")
            .zip(library_paths.into_iter())
            .flat_map(|x| vec![x.0, x.1])
            .collect();

        let temp_dir = TempDir::new("mdbook")?;

        let src_dir = self.source_dir();
        let preprocess_context = PreprocessorContext {
            src_dir
        };

        LinkPreprocessor::new().run(&preprocess_context, &mut self.book)?;

        for item in self.iter() {
            if let BookItem::Chapter(ref ch) = *item {
                if !ch.path.as_os_str().is_empty() {
                    let path = self.source_dir().join(&ch.path);
                    let content = utils::fs::file_to_string(&path)?;
                    println!("[*]: Testing file: {:?}", path);

                    // write preprocessed file to tempdir
                    let path = temp_dir.path().join(&ch.path);
                    let mut tmpf = utils::fs::create_file(&path)?;
                    tmpf.write_all(content.as_bytes())?;

                    let output = Command::new("rustdoc")
                        .arg(&path)
                        .arg("--test")
                        .args(&library_args)
                        .output()?;

                    if !output.status.success() {
                        bail!(ErrorKind::Subprocess(
                            "Rustdoc returned an error".to_string(),
                            output
                        ));
                    }
                }
            }
        }
        Ok(())
    }

    /// The logic for determining where a backend should put its build
    /// artefacts.
    ///
    /// If there is only 1 renderer, put it in the directory pointed to by the
    /// `build.build_dir` key in `Config`. If there is more than one then the
    /// renderer gets its own directory within the main build dir.
    ///
    /// i.e. If there were only one renderer (in this case, the HTML renderer):
    ///
    /// - build/
    ///   - index.html
    ///   - ...
    ///
    /// Otherwise if there are multiple:
    ///
    /// - build/
    ///   - epub/
    ///     - my_awesome_book.epub
    ///   - html/
    ///     - index.html
    ///     - ...
    ///   - latex/
    ///     - my_awesome_book.tex
    ///
    pub fn build_dir_for(&self, backend_name: &str) -> PathBuf {
        let build_dir = self.root.join(&self.config.build.build_dir);

        if self.renderers.len() <= 1 {
            build_dir
        } else {
            build_dir.join(backend_name)
        }
    }

    /// Get the directory containing this book's source files.
    pub fn source_dir(&self) -> PathBuf {
        self.root.join(&self.config.book.src)
    }

    // FIXME: This really belongs as part of the `HtmlConfig`.
    #[doc(hidden)]
    pub fn theme_dir(&self) -> PathBuf {
        match self.config.html_config().and_then(|h| h.theme) {
            Some(d) => self.root.join(d),
            None => self.root.join("theme"),
        }
    }
}

/// Look at the `Config` and try to figure out what renderers to use.
fn determine_renderers(config: &Config) -> Vec<Box<Renderer>> {
    let mut renderers: Vec<Box<Renderer>> = Vec::new();

    if let Some(output_table) = config.get("output").and_then(|o| o.as_table()) {
        for (key, table) in output_table.iter() {
            // the "html" backend has its own Renderer
            if key == "html" {
                renderers.push(Box::new(HtmlHandlebars::new()));
            } else {
                let renderer = interpret_custom_renderer(key, table);
                renderers.push(renderer);
            }
        }
    }

    // if we couldn't find anything, add the HTML renderer as a default
    if renderers.is_empty() {
        renderers.push(Box::new(HtmlHandlebars::new()));
    }

    renderers
}

/// Look at the `MDBook` and try to figure out what preprocessors to run.
fn determine_preprocessors(config: &Config) -> Result<Vec<Box<Preprocessor>>> {

    let preprocess_list = match config.build.preprocess {
        Some(ref p) => p,
        // If no preprocessor field is set, default to the LinkPreprocessor. This allows you
        // to disable the LinkPreprocessor by setting "preprocess" to an empty list.
        None => return Ok(vec![Box::new(LinkPreprocessor::new())])
    };

    let mut preprocessors: Vec<Box<Preprocessor>> = Vec::new();

    for key in preprocess_list {
        if key == "links" {
            preprocessors.push(Box::new(LinkPreprocessor::new()))
        } else {
            bail!("{:?} is not a recognised preprocessor", key);
        }
    }

    Ok(preprocessors)
}

fn interpret_custom_renderer(key: &str, table: &Value) -> Box<Renderer> {
    // look for the `command` field, falling back to using the key
    // prepended by "mdbook-"
    let table_dot_command = table
        .get("command")
        .and_then(|c| c.as_str())
        .map(|s| s.to_string());

    let command = table_dot_command.unwrap_or_else(|| format!("mdbook-{}", key));

    Box::new(CmdRenderer::new(key.to_string(), command.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use toml::value::{Table, Value};

    #[test]
    fn config_defaults_to_html_renderer_if_empty() {
        let cfg = Config::default();

        // make sure we haven't got anything in the `output` table
        assert!(cfg.get("output").is_none());

        let got = determine_renderers(&cfg);

        assert_eq!(got.len(), 1);
        assert_eq!(got[0].name(), "html");
    }

    #[test]
    fn add_a_random_renderer_to_the_config() {
        let mut cfg = Config::default();
        cfg.set("output.random", Table::new()).unwrap();

        let got = determine_renderers(&cfg);

        assert_eq!(got.len(), 1);
        assert_eq!(got[0].name(), "random");
    }

    #[test]
    fn add_a_random_renderer_with_custom_command_to_the_config() {
        let mut cfg = Config::default();

        let mut table = Table::new();
        table.insert("command".to_string(), Value::String("false".to_string()));
        cfg.set("output.random", table).unwrap();

        let got = determine_renderers(&cfg);

        assert_eq!(got.len(), 1);
        assert_eq!(got[0].name(), "random");
    }

    #[test]
    fn config_defaults_to_link_preprocessor_if_not_set() {
        let cfg = Config::default();

        // make sure we haven't got anything in the `output` table
        assert!(cfg.build.preprocess.is_none());

        let got = determine_preprocessors(&cfg);

        assert!(got.is_ok());
        assert_eq!(got.as_ref().unwrap().len(), 1);
        assert_eq!(got.as_ref().unwrap()[0].name(), "links");
    }

    #[test]
    fn config_doesnt_default_if_empty() {
        let cfg_str: &'static str = r#"
        [book]
        title = "Some Book"

        [build]
        build-dir = "outputs"
        create-missing = false
        preprocess = []
        "#;

        let cfg = Config::from_str(cfg_str).unwrap();

        // make sure we have something in the `output` table
        assert!(cfg.build.preprocess.is_some());

        let got = determine_preprocessors(&cfg);

        assert!(got.is_ok());
        assert!(got.unwrap().is_empty());
    }

    #[test]
    fn config_complains_if_unimplemented_preprocessor() {
        let cfg_str: &'static str = r#"
        [book]
        title = "Some Book"

        [build]
        build-dir = "outputs"
        create-missing = false
        preprocess = ["random"]
        "#;

        let cfg = Config::from_str(cfg_str).unwrap();

        // make sure we have something in the `output` table
        assert!(cfg.build.preprocess.is_some());

        let got = determine_preprocessors(&cfg);

        assert!(got.is_err());
    }
}
<commit_msg>Pull out default preprocessors into function<commit_after>//! The internal representation of a book and infrastructure for loading it from
//! disk and building it.
//!
//! For examples on using `MDBook`, consult the [top-level documentation][1].
//!
//! [1]: ../index.html

#![deny(missing_docs)]

mod summary;
mod book;
mod init;

pub use self::book::{load_book, Book, BookItem, BookItems, Chapter};
pub use self::summary::{parse_summary, Link, SectionNumber, Summary, SummaryItem};
pub use self::init::BookBuilder;

use std::path::PathBuf;
use std::io::Write;
use std::process::Command;
use tempdir::TempDir;
use toml::Value;

use utils;
use renderer::{CmdRenderer, HtmlHandlebars, RenderContext, Renderer};
use preprocess::{Preprocessor, LinkPreprocessor, PreprocessorContext};
use errors::*;

use config::Config;

/// The object used to manage and build a book.
pub struct MDBook {
    /// The book's root directory.
    pub root: PathBuf,
    /// The configuration used to tweak now a book is built.
    pub config: Config,
    /// A representation of the book's contents in memory.
    pub book: Book,
    renderers: Vec<Box<Renderer>>,

    /// The URL used for live reloading when serving up the book.
    pub livereload: Option<String>,

    /// List of pre-processors to be run on the book
    preprocessors: Vec<Box<Preprocessor>>
}

impl MDBook {
    /// Load a book from its root directory on disk.
    pub fn load<P: Into<PathBuf>>(book_root: P) -> Result<MDBook> {
        let book_root = book_root.into();
        let config_location = book_root.join("book.toml");

        // the book.json file is no longer used, so we should emit a warning to
        // let people know to migrate to book.toml
        if book_root.join("book.json").exists() {
            warn!("It appears you are still using book.json for configuration.");
            warn!("This format is no longer used, so you should migrate to the");
            warn!("book.toml format.");
            warn!("Check the user guide for migration information:");
            warn!("\thttps://rust-lang-nursery.github.io/mdBook/format/config.html");
        }

        let mut config = if config_location.exists() {
            debug!("[*] Loading config from {}", config_location.display());
            Config::from_disk(&config_location)?
        } else {
            Config::default()
        };

        config.update_from_env();

        if log_enabled!(::log::Level::Trace) {
            for line in format!("Config: {:#?}", config).lines() {
                trace!("{}", line);
            }
        }

        MDBook::load_with_config(book_root, config)
    }

    /// Load a book from its root directory using a custom config.
    pub fn load_with_config<P: Into<PathBuf>>(book_root: P, config: Config) -> Result<MDBook> {
        let root = book_root.into();

        let src_dir = root.join(&config.book.src);
        let book = book::load_book(&src_dir, &config.build)?;
        let livereload = None;

        let renderers = determine_renderers(&config);
        let preprocessors = determine_preprocessors(&config)?;

        Ok(MDBook {
            root,
            config,
            book,
            renderers,
            livereload,
            preprocessors,
        })
    }

    /// 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::book::BookItem;
    /// # #[allow(unused_variables)]
    /// # fn main() {
    /// # let book = MDBook::load("mybook").unwrap();
    /// for item in book.iter() {
    ///     match *item {
    ///         BookItem::Chapter(ref chapter) => {},
    ///         BookItem::Separator => {},
    ///     }
    /// }
    ///
    /// // 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 {
        self.book.iter()
    }

    /// `init()` gives you a `BookBuilder` which you can use to setup a new book
    /// and its accompanying directory structure.
    ///
    /// The `BookBuilder` 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 path provided as the root directory for your book, then adds
    /// in a `src/` directory containing a `SUMMARY.md` and `chapter_1.md` file
    /// to get you started.
    pub fn init<P: Into<PathBuf>>(book_root: P) -> BookBuilder {
        BookBuilder::new(book_root)
    }

    /// Tells the renderer to build our book and put it in the build directory.
    pub fn build(&self) -> Result<()> {
        debug!("[fn]: build");

        let mut preprocessed_book = self.book.clone();
        let preprocess_ctx = PreprocessorContext {
            src_dir: self.source_dir(),
        };

        for preprocessor in &self.preprocessors {
            preprocessor.run(&preprocess_ctx, &mut preprocessed_book)?;
        }

        for renderer in &self.renderers {
            self.run_renderer(&preprocessed_book, renderer.as_ref())?;
        }

        Ok(())
    }

    fn run_renderer(&self, preprocessed_book: &Book, renderer: &Renderer) -> Result<()> {
        let name = renderer.name();
        let build_dir = self.build_dir_for(name);
        if build_dir.exists() {
            debug!(
                "Cleaning build dir for the \"{}\" renderer ({})",
                name,
                build_dir.display()
            );

            utils::fs::remove_dir_content(&build_dir)
                .chain_err(|| "Unable to clear output directory")?;
        }

        let render_context = RenderContext::new(
            self.root.clone(),
            preprocessed_book.clone(),
            self.config.clone(),
            build_dir,
        );

        renderer
            .render(&render_context)
            .chain_err(|| "Rendering failed")
    }

    /// 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)
    pub fn with_renderer<R: Renderer + 'static>(&mut self, renderer: R) -> &mut Self {
        self.renderers.push(Box::new(renderer));
        self
    }

    /// You can add a new preprocessor by using this method.
    /// The only requirement is for your renderer to implement the Preprocessor trait.
    pub fn with_preprecessor<P: Preprocessor + 'static>(&mut self, preprocessor: P) -> &mut Self {
        self.preprocessors.push(Box::new(preprocessor));
        self
    }

    /// Run `rustdoc` tests on the book, linking against the provided libraries.
    pub fn test(&mut self, library_paths: Vec<&str>) -> Result<()> {
        let library_args: Vec<&str> = (0..library_paths.len())
            .map(|_| "-L")
            .zip(library_paths.into_iter())
            .flat_map(|x| vec![x.0, x.1])
            .collect();

        let temp_dir = TempDir::new("mdbook")?;

        let src_dir = self.source_dir();
        let preprocess_context = PreprocessorContext {
            src_dir
        };

        LinkPreprocessor::new().run(&preprocess_context, &mut self.book)?;

        for item in self.iter() {
            if let BookItem::Chapter(ref ch) = *item {
                if !ch.path.as_os_str().is_empty() {
                    let path = self.source_dir().join(&ch.path);
                    let content = utils::fs::file_to_string(&path)?;
                    println!("[*]: Testing file: {:?}", path);

                    // write preprocessed file to tempdir
                    let path = temp_dir.path().join(&ch.path);
                    let mut tmpf = utils::fs::create_file(&path)?;
                    tmpf.write_all(content.as_bytes())?;

                    let output = Command::new("rustdoc")
                        .arg(&path)
                        .arg("--test")
                        .args(&library_args)
                        .output()?;

                    if !output.status.success() {
                        bail!(ErrorKind::Subprocess(
                            "Rustdoc returned an error".to_string(),
                            output
                        ));
                    }
                }
            }
        }
        Ok(())
    }

    /// The logic for determining where a backend should put its build
    /// artefacts.
    ///
    /// If there is only 1 renderer, put it in the directory pointed to by the
    /// `build.build_dir` key in `Config`. If there is more than one then the
    /// renderer gets its own directory within the main build dir.
    ///
    /// i.e. If there were only one renderer (in this case, the HTML renderer):
    ///
    /// - build/
    ///   - index.html
    ///   - ...
    ///
    /// Otherwise if there are multiple:
    ///
    /// - build/
    ///   - epub/
    ///     - my_awesome_book.epub
    ///   - html/
    ///     - index.html
    ///     - ...
    ///   - latex/
    ///     - my_awesome_book.tex
    ///
    pub fn build_dir_for(&self, backend_name: &str) -> PathBuf {
        let build_dir = self.root.join(&self.config.build.build_dir);

        if self.renderers.len() <= 1 {
            build_dir
        } else {
            build_dir.join(backend_name)
        }
    }

    /// Get the directory containing this book's source files.
    pub fn source_dir(&self) -> PathBuf {
        self.root.join(&self.config.book.src)
    }

    // FIXME: This really belongs as part of the `HtmlConfig`.
    #[doc(hidden)]
    pub fn theme_dir(&self) -> PathBuf {
        match self.config.html_config().and_then(|h| h.theme) {
            Some(d) => self.root.join(d),
            None => self.root.join("theme"),
        }
    }
}

/// Look at the `Config` and try to figure out what renderers to use.
fn determine_renderers(config: &Config) -> Vec<Box<Renderer>> {
    let mut renderers: Vec<Box<Renderer>> = Vec::new();

    if let Some(output_table) = config.get("output").and_then(|o| o.as_table()) {
        for (key, table) in output_table.iter() {
            // the "html" backend has its own Renderer
            if key == "html" {
                renderers.push(Box::new(HtmlHandlebars::new()));
            } else {
                let renderer = interpret_custom_renderer(key, table);
                renderers.push(renderer);
            }
        }
    }

    // if we couldn't find anything, add the HTML renderer as a default
    if renderers.is_empty() {
        renderers.push(Box::new(HtmlHandlebars::new()));
    }

    renderers
}

fn default_preprocessors() -> Vec<Box<Preprocessor>> {
    vec![Box::new(LinkPreprocessor::new())]
}

/// Look at the `MDBook` and try to figure out what preprocessors to run.
fn determine_preprocessors(config: &Config) -> Result<Vec<Box<Preprocessor>>> {

    let preprocess_list = match config.build.preprocess {
        Some(ref p) => p,
        // If no preprocessor field is set, default to the LinkPreprocessor. This allows you
        // to disable the LinkPreprocessor by setting "preprocess" to an empty list.
        None => return Ok(default_preprocessors())
    };

    let mut preprocessors: Vec<Box<Preprocessor>> = Vec::new();

    for key in preprocess_list {
        if key == "links" {
            preprocessors.push(Box::new(LinkPreprocessor::new()))
        } else {
            bail!("{:?} is not a recognised preprocessor", key);
        }
    }

    Ok(preprocessors)
}

fn interpret_custom_renderer(key: &str, table: &Value) -> Box<Renderer> {
    // look for the `command` field, falling back to using the key
    // prepended by "mdbook-"
    let table_dot_command = table
        .get("command")
        .and_then(|c| c.as_str())
        .map(|s| s.to_string());

    let command = table_dot_command.unwrap_or_else(|| format!("mdbook-{}", key));

    Box::new(CmdRenderer::new(key.to_string(), command.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use toml::value::{Table, Value};

    #[test]
    fn config_defaults_to_html_renderer_if_empty() {
        let cfg = Config::default();

        // make sure we haven't got anything in the `output` table
        assert!(cfg.get("output").is_none());

        let got = determine_renderers(&cfg);

        assert_eq!(got.len(), 1);
        assert_eq!(got[0].name(), "html");
    }

    #[test]
    fn add_a_random_renderer_to_the_config() {
        let mut cfg = Config::default();
        cfg.set("output.random", Table::new()).unwrap();

        let got = determine_renderers(&cfg);

        assert_eq!(got.len(), 1);
        assert_eq!(got[0].name(), "random");
    }

    #[test]
    fn add_a_random_renderer_with_custom_command_to_the_config() {
        let mut cfg = Config::default();

        let mut table = Table::new();
        table.insert("command".to_string(), Value::String("false".to_string()));
        cfg.set("output.random", table).unwrap();

        let got = determine_renderers(&cfg);

        assert_eq!(got.len(), 1);
        assert_eq!(got[0].name(), "random");
    }

    #[test]
    fn config_defaults_to_link_preprocessor_if_not_set() {
        let cfg = Config::default();

        // make sure we haven't got anything in the `output` table
        assert!(cfg.build.preprocess.is_none());

        let got = determine_preprocessors(&cfg);

        assert!(got.is_ok());
        assert_eq!(got.as_ref().unwrap().len(), 1);
        assert_eq!(got.as_ref().unwrap()[0].name(), "links");
    }

    #[test]
    fn config_doesnt_default_if_empty() {
        let cfg_str: &'static str = r#"
        [book]
        title = "Some Book"

        [build]
        build-dir = "outputs"
        create-missing = false
        preprocess = []
        "#;

        let cfg = Config::from_str(cfg_str).unwrap();

        // make sure we have something in the `output` table
        assert!(cfg.build.preprocess.is_some());

        let got = determine_preprocessors(&cfg);

        assert!(got.is_ok());
        assert!(got.unwrap().is_empty());
    }

    #[test]
    fn config_complains_if_unimplemented_preprocessor() {
        let cfg_str: &'static str = r#"
        [book]
        title = "Some Book"

        [build]
        build-dir = "outputs"
        create-missing = false
        preprocess = ["random"]
        "#;

        let cfg = Config::from_str(cfg_str).unwrap();

        // make sure we have something in the `output` table
        assert!(cfg.build.preprocess.is_some());

        let got = determine_preprocessors(&cfg);

        assert!(got.is_err());
    }
}
<|endoftext|>


<commit_before>use adler32::RollingAdler32;

pub fn adler32_from_bytes(bytes: &[u8; 4]) -> u32 {
    let val: u32 = (bytes[0] as u32) | ((bytes[1] as u32) << 8) |
       ((bytes[2] as u32) << 16) | ((bytes[3] as u32) << 24);
    u32::from_be(val)
}

/// Whether we should validate the checksum, and what type of checksum it is.
pub enum ChecksumType {
    /// No validation.
    ///
    /// For raw deflate streams or when we don't bother checking.
    None,
    /// Adler32
    ///
    /// Used in the zlib format.
    Adler32(RollingAdler32),
}

pub struct Checksum {
    checksum_type: ChecksumType,
}

impl Checksum {
    #[inline]
    pub fn none() -> Checksum {
        Checksum::new(ChecksumType::None)
    }

    #[inline]
    pub fn is_none(&self) -> bool {
        match self.checksum_type {
            ChecksumType::None => true,
            _ => false,
        }
    }

    #[inline]
    pub fn zlib() -> Checksum {
        Checksum::new(ChecksumType::Adler32(RollingAdler32::new()))
    }

    pub fn new(checksum_type: ChecksumType) -> Checksum {
        Checksum {
            checksum_type: checksum_type,
        }
    }

    #[inline]
    pub fn update(&mut self, bytes: &[u8]) {
        match self.checksum_type {
            ChecksumType::None => (),
            ChecksumType::Adler32(ref mut c) => {
                c.update_buffer(bytes);
            }
        }
    }

    pub fn check(&self, expected: u32) -> Result<(), String> {
        match self.checksum_type {
            ChecksumType::None => Ok(()),
            ChecksumType::Adler32(ref c) => {
                if c.hash() == expected {
                    Ok(())
                } else {
                    Err("Checksum mismatch!".to_owned())
                }
            },
        }

    }

    #[inline]
    pub fn current_value(&self) -> u32 {
        match self.checksum_type {
            ChecksumType::Adler32(ref c) => c.hash(),
            _ => 0,
        }
    }

}

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

    #[test]
    fn adler32() {
        let bytes = [0x00, 0x00, 0x01, 0x0b];
        assert_eq!(adler32_from_bytes(&bytes), 267);
    }
}
<commit_msg>Fix Checksum Endianness<commit_after>use adler32::RollingAdler32;

pub fn adler32_from_bytes(bytes: &[u8; 4]) -> u32 {
    (bytes[3] as u32) | ((bytes[2] as u32) << 8) |
       ((bytes[1] as u32) << 16) | ((bytes[0] as u32) << 24)
}

/// Whether we should validate the checksum, and what type of checksum it is.
pub enum ChecksumType {
    /// No validation.
    ///
    /// For raw deflate streams or when we don't bother checking.
    None,
    /// Adler32
    ///
    /// Used in the zlib format.
    Adler32(RollingAdler32),
}

pub struct Checksum {
    checksum_type: ChecksumType,
}

impl Checksum {
    #[inline]
    pub fn none() -> Checksum {
        Checksum::new(ChecksumType::None)
    }

    #[inline]
    pub fn is_none(&self) -> bool {
        match self.checksum_type {
            ChecksumType::None => true,
            _ => false,
        }
    }

    #[inline]
    pub fn zlib() -> Checksum {
        Checksum::new(ChecksumType::Adler32(RollingAdler32::new()))
    }

    pub fn new(checksum_type: ChecksumType) -> Checksum {
        Checksum {
            checksum_type: checksum_type,
        }
    }

    #[inline]
    pub fn update(&mut self, bytes: &[u8]) {
        match self.checksum_type {
            ChecksumType::None => (),
            ChecksumType::Adler32(ref mut c) => {
                c.update_buffer(bytes);
            }
        }
    }

    pub fn check(&self, expected: u32) -> Result<(), String> {
        match self.checksum_type {
            ChecksumType::None => Ok(()),
            ChecksumType::Adler32(ref c) => {
                if c.hash() == expected {
                    Ok(())
                } else {
                    Err("Checksum mismatch!".to_owned())
                }
            },
        }

    }

    #[inline]
    pub fn current_value(&self) -> u32 {
        match self.checksum_type {
            ChecksumType::Adler32(ref c) => c.hash(),
            _ => 0,
        }
    }

}

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

    #[test]
    fn adler32() {
        let bytes = [0x00, 0x00, 0x01, 0x0b];
        assert_eq!(adler32_from_bytes(&bytes), 267);
    }
}
<|endoftext|>


<commit_before>//Copyright (c) 2016, Ruslan Baratov, Alex Frappier Lachapelle
//All rights reserved.

use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};

use config;
use webhooks;

include!("logger_macros.rs");

////////////////////////////////////////////////////////////
//                         Command                        //
////////////////////////////////////////////////////////////

//                            Config,                             Raw webhook,            Arguments
pub type CommandCallback = fn(&Arc<Mutex<config::ConfigHandler>>, webhooks::WebhookEvent, Vec<&str>) -> Result<String, String>;

pub struct Command {
    pub requires_please: bool,
    pub whitelist_only:  bool,
    pub callback:        CommandCallback
}

impl Command {
    pub fn new(requires_please: bool, whitelist_only: bool, callback: CommandCallback) -> Command {
        Command{
            requires_please: requires_please,
            whitelist_only:  whitelist_only,
            callback:        callback
        }
    }
}

////////////////////////////////////////////////////////////
//                     CommandHandler                     //
////////////////////////////////////////////////////////////

pub struct CommandHandler {
    config:   Arc<Mutex<config::ConfigHandler>>,
    commands: BTreeMap<String, Command>
}

impl CommandHandler {

    pub fn new(tsconfig: &Arc<Mutex<config::ConfigHandler>>) -> CommandHandler {

        //Register commands
        let mut commands: BTreeMap<String, Command> = BTreeMap::new();
        commands.insert(String::from("ping"), Command::new(false, false, ping));

        CommandHandler {
            config:   tsconfig.clone(),
            commands: commands
        }
    }

    pub fn parse_command(&self, webhook: webhooks::WebhookEvent) {

        let mut tokens: Vec<&str>   = webhook.command.split_whitespace().collect();
        let mut is_please_provided  = false;
        let mut is_user_whitelisted = false;
        let mut next_token_index    = 0;

        //Check if please was said
        if tokens[0].to_lowercase() == "please" {
            is_please_provided = true;
            next_token_index   = 1;
        }

        //Check if user is whitelisted for restricted commands
        //  Restricted scope, we are using the config here but
        //  the callback can also use it so we force the mutex
        //  to unlock before letting the callback lock it
        //  given that locking twice will induce a panic.
        {
            let mut config = self.config.lock().unwrap();
            is_user_whitelisted = config.whitelist_validate_user(webhook.clone().user);
        }

        //Find command among registered commands
        let command = match self.commands.get(tokens[next_token_index]) {
            Some(command) => command,
            None          => {
                respond(&self.config, webhook.clone(), String::from("Sorry the command was not found."));
                return;
            }
        };

        //Check if please and whitelist is required
        let mut run_cmd         = false;
        let mut response_prefix = String::new();
        if !(command.whitelist_only && !is_user_whitelisted) {
            if (command.requires_please && is_please_provided) || (!command.requires_please && !is_please_provided) {
                run_cmd = true;
            } else if !command.requires_please && is_please_provided {
                run_cmd         = true;
                response_prefix = String::from("You didn't need to say please but thanks anyways :smiley:\n");
            } else if command.requires_please && !is_please_provided {
                respond(&self.config, webhook.clone(), String::from("Whats the magic word?"));
            }
        } else {
            respond(&self.config, webhook.clone(), String::from("Sorry! That command if for whitelisted people only!"));
        }

        if run_cmd {
            match (command.callback)(&self.config, webhook.clone(), tokens.split_off(next_token_index)) {
                Ok(msg)  => {
                    respond(&self.config, webhook.clone(), format!("{}{}", response_prefix, msg));
                }
                Err(msg) => {
                    respond(&self.config, webhook.clone(), format!("An error occurred while executing the command: {}", msg));
                }
            }
        }
    }
}

////////////////////////////////////////////////////////////
//                        Callbacks                       //
////////////////////////////////////////////////////////////

pub fn ping(tsconfig: &Arc<Mutex<config::ConfigHandler>>, raw_event: webhooks::WebhookEvent, Args: Vec<&str>) -> Result<String, String> {
    return Ok(String::from("Pong"))
}

////////////////////////////////////////////////////////////
//                          Utils                         //
////////////////////////////////////////////////////////////

pub fn respond(tsconfig: &Arc<Mutex<config::ConfigHandler>>, raw_event: webhooks::WebhookEvent, msg: String) {

    //TODO: Don't forget to change this when refactoring config
    //Get repo were following
    let github_follow_repo: String;
    let github_owner_token: String;
    {
        let mut config = tsconfig.lock().unwrap();
        match config.get_string("config", "github_follow_repo") {
            Ok(_github_follow_repo) => github_follow_repo = _github_follow_repo,
            Err(err)                => {
                thread_error!("Failed to acquire \"github_follow_repo\": {}", err);
                panic!("Failed to acquire \"github_follow_repo\": {}", err);
            }
        }

        //Get owner api token
        match config.get_string("config", "github_owner_token") {
            Ok(_github_owner_token) => github_owner_token = _github_owner_token,
            Err(err)                => {
                thread_error!("Error getting  the \"github_owner_token\" value from config: {}", err);
                panic!("Error getting the \"github_owner_token\" value from config: {}", err);
            }
        }
    }

    let endpoint = format!("repos/{}/issues/{}/comments?access_token={}", github_follow_repo, raw_event.number, github_owner_token);
    let message  = format!("{{body: \"@{} {}\"}}", raw_event.user, msg);
    match webhooks::github_post_request(endpoint, message) {
        Ok(())   => (),
        Err(err) => {thread_error!("{}", err);}
    }
}
<commit_msg>Fixed response error<commit_after>//Copyright (c) 2016, Ruslan Baratov, Alex Frappier Lachapelle
//All rights reserved.

use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};

use config;
use webhooks;

include!("logger_macros.rs");

////////////////////////////////////////////////////////////
//                         Command                        //
////////////////////////////////////////////////////////////

//                            Config,                             Raw webhook,            Arguments
pub type CommandCallback = fn(&Arc<Mutex<config::ConfigHandler>>, webhooks::WebhookEvent, Vec<&str>) -> Result<String, String>;

pub struct Command {
    pub requires_please: bool,
    pub whitelist_only:  bool,
    pub callback:        CommandCallback
}

impl Command {
    pub fn new(requires_please: bool, whitelist_only: bool, callback: CommandCallback) -> Command {
        Command{
            requires_please: requires_please,
            whitelist_only:  whitelist_only,
            callback:        callback
        }
    }
}

////////////////////////////////////////////////////////////
//                     CommandHandler                     //
////////////////////////////////////////////////////////////

pub struct CommandHandler {
    config:   Arc<Mutex<config::ConfigHandler>>,
    commands: BTreeMap<String, Command>
}

impl CommandHandler {

    pub fn new(tsconfig: &Arc<Mutex<config::ConfigHandler>>) -> CommandHandler {

        //Register commands
        let mut commands: BTreeMap<String, Command> = BTreeMap::new();
        commands.insert(String::from("ping"), Command::new(false, false, ping));

        CommandHandler {
            config:   tsconfig.clone(),
            commands: commands
        }
    }

    pub fn parse_command(&self, webhook: webhooks::WebhookEvent) {

        let mut tokens: Vec<&str>   = webhook.command.split_whitespace().collect();
        let mut is_please_provided  = false;
        let mut is_user_whitelisted = false;
        let mut next_token_index    = 0;

        //Check if please was said
        if tokens[0].to_lowercase() == "please" {
            is_please_provided = true;
            next_token_index   = 1;
        }

        //Check if user is whitelisted for restricted commands
        //  Restricted scope, we are using the config here but
        //  the callback can also use it so we force the mutex
        //  to unlock before letting the callback lock it
        //  given that locking twice will induce a panic.
        {
            let mut config = self.config.lock().unwrap();
            is_user_whitelisted = config.whitelist_validate_user(webhook.clone().user);
        }

        //Find command among registered commands
        let command = match self.commands.get(tokens[next_token_index]) {
            Some(command) => command,
            None          => {
                respond(&self.config, webhook.clone(), String::from("Sorry the command was not found."));
                return;
            }
        };

        //Check if please and whitelist is required
        let mut run_cmd         = false;
        let mut response_prefix = String::new();
        if !(command.whitelist_only && !is_user_whitelisted) {
            if (command.requires_please && is_please_provided) || (!command.requires_please && !is_please_provided) {
                run_cmd = true;
            } else if !command.requires_please && is_please_provided {
                run_cmd         = true;
                response_prefix = String::from("You didn't need to say please but thanks anyways :smiley:\n");
            } else if command.requires_please && !is_please_provided {
                respond(&self.config, webhook.clone(), String::from("Whats the magic word?"));
            }
        } else {
            respond(&self.config, webhook.clone(), String::from("Sorry! That command if for whitelisted people only!"));
        }

        if run_cmd {
            match (command.callback)(&self.config, webhook.clone(), tokens.split_off(next_token_index)) {
                Ok(msg)  => {
                    respond(&self.config, webhook.clone(), format!("{}{}", response_prefix, msg));
                }
                Err(msg) => {
                    respond(&self.config, webhook.clone(), format!("An error occurred while executing the command: {}", msg));
                }
            }
        }
    }
}

////////////////////////////////////////////////////////////
//                        Callbacks                       //
////////////////////////////////////////////////////////////

pub fn ping(tsconfig: &Arc<Mutex<config::ConfigHandler>>, raw_event: webhooks::WebhookEvent, Args: Vec<&str>) -> Result<String, String> {
    return Ok(String::from("Pong"))
}

////////////////////////////////////////////////////////////
//                          Utils                         //
////////////////////////////////////////////////////////////

pub fn respond(tsconfig: &Arc<Mutex<config::ConfigHandler>>, raw_event: webhooks::WebhookEvent, msg: String) {

    //TODO: Don't forget to change this when refactoring config
    //Get repo were following
    let github_follow_repo: String;
    let github_owner_token: String;
    {
        let mut config = tsconfig.lock().unwrap();
        match config.get_string("config", "github_follow_repo") {
            Ok(_github_follow_repo) => github_follow_repo = _github_follow_repo,
            Err(err)                => {
                thread_error!("Failed to acquire \"github_follow_repo\": {}", err);
                panic!("Failed to acquire \"github_follow_repo\": {}", err);
            }
        }

        //Get owner api token
        match config.get_string("config", "github_bot_token") {
            Ok(_github_owner_token) => github_owner_token = _github_owner_token,
            Err(err)                => {
                thread_error!("Error getting  the \"github_owner_token\" value from config: {}", err);
                panic!("Error getting the \"github_owner_token\" value from config: {}", err);
            }
        }
    }

    let endpoint = format!("repos/{}/issues/{}/comments?access_token={}", github_follow_repo, raw_event.number, github_owner_token);
    let message  = format!("{{\"body\": \"@{} {}\"}}", raw_event.user, msg);
    match webhooks::github_post_request(endpoint, message) {
        Ok(())   => (),
        Err(err) => {thread_error!("{}", err);}
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;

use std::io::Cursor;
use self::byteorder::{ReadBytesExt, WriteBytesExt, BigEndian, LittleEndian};

pub struct Value {
    data: Vec<u8>
}

impl From<i64> for Value {
    fn from(val: i64) -> Value {
        let mut buffer = Vec::new();
        buffer.write_i64::<BigEndian>(val).unwrap();
        Value{data: buffer}
    }
}<commit_msg>String support<commit_after>extern crate byteorder;

use std::io::Cursor;
use self::byteorder::{ReadBytesExt, WriteBytesExt, BigEndian, LittleEndian};

pub struct Value {
    data: Vec<u8>
}

impl From<i64> for Value {
    fn from(val: i64) -> Value {
        let mut buffer = Vec::new();
        buffer.write_i64::<BigEndian>(val).unwrap();
        Value { data: buffer }
    }
}

impl From<String> for Value {
    fn from(val: String) -> Value {
        Value{data:val.into_bytes()}
    }
}<|endoftext|>


<commit_before>use {Byte, Chunk, SChunks};
use client::GetResponse;
use hyper::client::Client;
use hyper::error::Error;
use hyper::header::{ByteRangeSpec, Headers, Range};
use pbr::{MultiBar, Pipe, ProgressBar};
use std::cmp::min;
use std::io::Read;
use std::thread;

/// Represents a range between two Byte types
#[derive(Debug, PartialEq)]
struct RangeBytes(Byte, Byte);

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Byte,
                    global_chunk_length: Byte)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Byte = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Byte = min(content_length - 1,
                            ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

/// Function to get the HTTP header to send to the file server, for a chunk (specified by its index)
fn get_header_from_index(chunk_index: u64,
                         content_length: Byte,
                         global_chunk_length: Byte)
                         -> Option<(Headers, Byte)> {

    match get_chunk_length(chunk_index, content_length, global_chunk_length) {
        Some(range) => {
            let mut header = Headers::new();
            header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(range.0, range.1)]));
            Some((header, range.1 - range.0))
        }
        None => None,
    }

}

/// Function to get from the server the content of a chunk.
/// This function returns a Result type - Byte if the content of the header is accessible, an Error type otherwise.
fn download_a_chunk(http_client: &Client,
                    http_header: Headers,
                    chunk_vector: &mut Chunk,
                    url: &str,
                    mpb: &mut ProgressBar<Pipe>)
                    -> Result<Byte, Error> {

    match http_client.get_http_response_using_headers(url, http_header) {
        Ok(mut body) => {
            let mut bytes_buffer = [0; 2048];
            let mut sum_bytes = 0;
            while let Ok(n) = body.read(&mut bytes_buffer) {
                if n == 0 {
                    return Ok(sum_bytes);
                }
                chunk_vector.extend_from_slice(&bytes_buffer[..n]);
                sum_bytes += n as u64;
                mpb.add(n as u64);
            }
            return Ok(0u64);
        }
        Err(http_error) => Err(http_error),
    }

}

/// Function to download each chunk of a remote content (given by its URL).
/// This function takes as parameters:
/// * the remote content length,
/// * a mutable reference to share between threads, which contains each chunk,
/// * the number of chunks that contains the remote content,
/// * the URL of the remote content server.
pub fn download_chunks(content_length: u64,
                       downloaded_chunks: &SChunks,
                       nb_chunks: u64,
                       url: &str) {

    // let mut downloaded_chunks: Arc<Mutex<Chunks>> =
    //     Arc::new(Mutex::new(Chunks::with_capacity(nb_chunks as usize)));
    let global_chunk_length: u64 = (content_length / nb_chunks) + 1;
    let mut jobs = vec![];

    let mut mpb = MultiBar::new();
    mpb.println(&format!("Downloading {} chunks: ", nb_chunks));

    for chunk_index in 0..nb_chunks {

        let (http_header, chunk_length) =
            get_header_from_index(chunk_index, content_length, global_chunk_length).unwrap();
        let mut chunk_content = Chunk::with_capacity(chunk_length as usize);
        let hyper_client = Client::new();
        let url_clone = String::from(url);
        let clone_chunks = downloaded_chunks.clone();

        let mut mp = mpb.create_bar(chunk_length);
        mp.message(&format!("Chunk {} - ", chunk_index));

        jobs.push(thread::spawn(move || {

            match download_a_chunk(&hyper_client,
                                   http_header,
                                   &mut chunk_content,
                                   &url_clone,
                                   &mut mp) {
                Ok(bytes_written) => {
                    if bytes_written > 0 {
                        let mut shared_clone_chunks = clone_chunks.lock().unwrap();
                        shared_clone_chunks.insert(chunk_index as usize, chunk_content);
                        mp.finish();
                    } else {
                        panic!("The downloaded chunk {} is empty", chunk_index);
                    }
                }
                Err(error) => {
                    panic!("Canno't download the chunk {}, due to error {}",
                           chunk_index,
                           error);
                }
            }
        }));
    }

    mpb.listen();

    for child in jobs {
        let _ = child.join();
    }

}

#[cfg(test)]
mod test_chunk_length {

    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}

#[cfg(test)]
mod test_header {

    use super::get_header_from_index;
    use hyper::header::{ByteRangeSpec, Headers, Range};

    #[test]
    fn wrong_chunk_length_should_return_none() {
        assert_eq!(None, get_header_from_index(0, 0, 0));
    }

    #[test]
    fn good_chunk_length_should_return_a_good_header() {
        let mut test_header = Headers::new();
        test_header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(750, 997)]));
        assert_eq!(Some((test_header, 247)), get_header_from_index(3, 998, 250));
    }

}
<commit_msg>Refactor code & update Byte type to Bytes<commit_after>use {Bytes, Chunk, SChunks};
use client::GetResponse;
use hyper::client::Client;
use hyper::error::Error;
use hyper::header::{ByteRangeSpec, Headers, Range};
use pbr::{MultiBar, Pipe, ProgressBar};
use std::cmp::min;
use std::io::Read;
use std::thread;

/// Represents a range between two Bytes types
#[derive(Debug, PartialEq)]
struct RangeBytes(Bytes, Bytes);

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Bytes,
                    global_chunk_length: Bytes)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Bytes = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Bytes = min(content_length - 1,
                             ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

/// Function to get the HTTP header to send to the file server, for a chunk (specified by its index)
fn get_header_from_index(chunk_index: u64,
                         content_length: Bytes,
                         global_chunk_length: Bytes)
                         -> Option<(Headers, Bytes)> {

    match get_chunk_length(chunk_index, content_length, global_chunk_length) {
        Some(range) => {
            let mut header = Headers::new();
            header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(range.0, range.1)]));
            Some((header, range.1 - range.0))
        }
        None => None,
    }

}

/// Function to get from the server the content of a chunk.
/// This function returns a Result type - Bytes if the content of the header is accessible, an Error type otherwise.
fn download_a_chunk(http_client: &Client,
                    http_header: Headers,
                    chunk_vector: &mut Chunk,
                    url: &str,
                    mpb: &mut ProgressBar<Pipe>)
                    -> Result<Bytes, Error> {

    match http_client.get_http_response_using_headers(url, http_header) {
        Ok(mut body) => {
            let mut bytes_buffer = [0; 2048];
            let mut sum_bytes = 0;
            while let Ok(n) = body.read(&mut bytes_buffer) {
                if n == 0 {
                    return Ok(sum_bytes);
                }
                chunk_vector.extend_from_slice(&bytes_buffer[..n]);
                sum_bytes += n as u64;
                mpb.add(n as u64);
            }
            return Ok(0u64);
        }
        Err(http_error) => Err(http_error),
    }

}

/// Function to download each chunk of a remote content (given by its URL).
/// This function takes as parameters:
/// * the remote content length,
/// * a mutable reference to share between threads, which contains each chunk,
/// * the number of chunks that contains the remote content,
/// * the URL of the remote content server.
pub fn download_chunks(content_length: u64,
                       downloaded_chunks: &SChunks,
                       nb_chunks: u64,
                       url: &str) {

    // let mut downloaded_chunks: Arc<Mutex<Chunks>> =
    //     Arc::new(Mutex::new(Chunks::with_capacity(nb_chunks as usize)));
    let global_chunk_length: u64 = (content_length / nb_chunks) + 1;
    let mut jobs = vec![];

    let mut mpb = MultiBar::new();
    mpb.println(&format!("Downloading {} chunks: ", nb_chunks));

    for chunk_index in 0..nb_chunks {

        let (http_header, chunk_length) =
            get_header_from_index(chunk_index, content_length, global_chunk_length).unwrap();
        let mut chunk_content = Chunk::with_capacity(chunk_length as usize);
        let hyper_client = Client::new();
        let url_clone = String::from(url);
        let clone_chunks = downloaded_chunks.clone();

        let mut mp = mpb.create_bar(chunk_length);
        mp.message(&format!("Chunk {} - ", chunk_index));

        jobs.push(thread::spawn(move || match download_a_chunk(&hyper_client,
                                                                    http_header,
                                                                    &mut chunk_content,
                                                                    &url_clone,
                                                                    &mut mp) {
            Ok(bytes_written) => {
                if bytes_written > 0 {
                    let mut shared_clone_chunks = clone_chunks.lock().unwrap();
                    shared_clone_chunks.insert(chunk_index as usize, chunk_content);
                    mp.finish();
                } else {
                    panic!("The downloaded chunk {} is empty", chunk_index);
                }
            }
            Err(error) => {
                panic!("Canno't download the chunk {}, due to error {}",
                       chunk_index,
                       error);
            }
        }));
    }

    mpb.listen();

    for child in jobs {
        let _ = child.join();
    }

}

#[cfg(test)]
mod test_chunk_length {

    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}

#[cfg(test)]
mod test_header {

    use super::get_header_from_index;
    use hyper::header::{ByteRangeSpec, Headers, Range};

    #[test]
    fn wrong_chunk_length_should_return_none() {
        assert_eq!(None, get_header_from_index(0, 0, 0));
    }

    #[test]
    fn good_chunk_length_should_return_a_good_header() {
        let mut test_header = Headers::new();
        test_header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(750, 997)]));
        assert_eq!(Some((test_header, 247)), get_header_from_index(3, 998, 250));
    }

}
<|endoftext|>


<commit_before>pub mod raw {
   #![allow(non_camel_case_types)]

   use libc::{c_char, c_uchar, c_short, c_ushort, c_int, c_uint, c_long, c_ulong};

   // run this hook first, before ANYTHING
   pub const APR_HOOK_REALLY_FIRST:  c_int = -10;
   // run this hook first
   pub const APR_HOOK_FIRST:         c_int = 0;
   // run this hook somewhere
   pub const APR_HOOK_MIDDLE:        c_int = 10;
   // run this hook after every other hook which is defined
   pub const APR_HOOK_LAST:          c_int = 20;
   // run this hook last, after EVERYTHING
   pub const APR_HOOK_REALLY_LAST:   c_int = 30;

   extern "C" {
      pub fn apr_table_get(t: *const apr_table_t, key: *const c_char) -> *const c_char;
      pub fn apr_table_set(t: *mut apr_table_t, key: *const c_char, val: *const c_char) -> ();
      pub fn apr_table_elts(t: *const apr_table_t) -> *const apr_array_header_t;
      pub fn apr_pstrmemdup(p: *mut apr_pool_t, s: *const c_char, n: apr_size_t) -> *mut c_char;
      pub fn apr_version_string() -> *const c_char;
      pub fn apu_version_string() -> *const c_char;
   }

   pub fn dup_c_str<T: Into<Vec<u8>>>(pool: *mut apr_pool_t, data: T) -> *mut c_char {
      let bytes = data.into();

      unsafe {
         apr_pstrmemdup(
            pool,
            bytes.as_ptr() as *const c_char,
            bytes.len() as apr_size_t
         )
      }
   }

   #[repr(C)]
   pub struct apr_bucket_alloc_t;

   #[repr(C)]
   pub struct apr_bucket_brigade;

   #[repr(C)]
   pub struct apr_finfo_t;

   #[repr(C)]
   pub struct apr_pool_t;

   #[repr(C)]
   pub struct apr_sockaddr_t;

   #[repr(C)]
   pub struct apr_table_t;

   #[repr(C)]
   pub struct apr_thread_mutex_t;

   #[repr(C)]
   pub struct apr_thread_t;

   #[repr(C)]
   pub struct apr_uri_t;

   pub type apr_byte_t = c_uchar;
   pub type apr_int16_t = c_short;
   pub type apr_uint16_t = c_ushort;
   pub type apr_int32_t = c_int;
   pub type apr_uint32_t = c_uint;
   pub type apr_int64_t = c_long;
   pub type apr_uint64_t = c_ulong;
   pub type apr_size_t = c_ulong;
   pub type apr_ssize_t = c_long;
   pub type apr_off_t = c_long;
   pub type apr_socklen_t = c_uint;
   pub type apr_ino_t = c_ulong;
   pub type apr_uintptr_t = apr_uint64_t;
   pub type apr_status_t = c_int;
   pub type apr_signum_t = c_int;
   pub type apr_time_t = apr_int64_t;
   pub type apr_port_t = apr_uint16_t;

   #[repr(C)]
   pub struct apr_array_header_t {
      pub pool: *mut apr_pool_t,
      pub elt_size: c_int,
      pub nelts: c_int,
      pub nalloc: c_int,
      pub elts: *mut c_char,
   }

   #[repr(C)]
   pub struct apr_table_entry_t {
      pub key: *mut c_char,
      pub val: *mut c_char,
      pub key_checksum: apr_uint32_t,
   }
}

use std::ffi::CString;

use wrapper::{Wrapper, c_str_value};

pub type AprTable<'a> = Wrapper<'a, raw::apr_table_t>;

impl<'a> AprTable<'a> {
   pub fn get(&self, key: &'static str) -> Option<&'a str> {
      let p: *const raw::apr_table_t = self.raw;
      c_str_value(
         unsafe { raw::apr_table_get(p, CString::new(key).unwrap().as_ptr()) }
      )
   }

   pub fn set(&mut self, key: &'static str, val: &'static str) {
      let p: *mut raw::apr_table_t = self.raw;
      unsafe {
         raw::apr_table_set(
            p,
            CString::new(key).unwrap().as_ptr(),
            CString::new(val).unwrap().as_ptr()
         )
      };
   }

   pub fn iter(&self) -> AprTableIter {
      let ptr = unsafe { raw::apr_table_elts(self.raw) };
      let raw: &raw::apr_array_header_t = unsafe { &*ptr };

      AprTableIter {
         array_header: raw,
         next_idx: 0
      }
   }
}

pub struct AprTableIter<'a> {
   array_header: &'a raw::apr_array_header_t,
   next_idx: usize,
}

impl<'a> Iterator for AprTableIter<'a> {
   type Item = (&'a str, &'a str);

   fn next(&mut self) -> Option<(&'a str, &'a str)> {
      if self.next_idx != self.array_header.nelts as usize {
         let mut elts = self.array_header.elts as *const raw::apr_table_entry_t;

         elts = unsafe { elts.offset(self.next_idx as isize) };
         self.next_idx += 1;

         let key = c_str_value(unsafe { (*elts).key }).unwrap();
         let val = c_str_value(unsafe { (*elts).val }).unwrap();

         Some((key, val))
      } else {
         None
      }
   }

   fn size_hint(&self) -> (usize, Option<usize>) {
      let rem = self.array_header.nelts as usize - self.next_idx;
      (rem, Some(rem))
   }
}

pub fn apr_version_string<'a>() -> Option<&'a str> {
   c_str_value(
      unsafe { raw::apr_version_string() }
   )
}

pub fn apu_version_string<'a>() -> Option<&'a str> {
   c_str_value(
      unsafe { raw::apu_version_string() }
   )
}
<commit_msg>No need to convert to raw pointers - use automatic conversion.<commit_after>pub mod raw {
   #![allow(non_camel_case_types)]

   use libc::{c_char, c_uchar, c_short, c_ushort, c_int, c_uint, c_long, c_ulong};

   // run this hook first, before ANYTHING
   pub const APR_HOOK_REALLY_FIRST:  c_int = -10;
   // run this hook first
   pub const APR_HOOK_FIRST:         c_int = 0;
   // run this hook somewhere
   pub const APR_HOOK_MIDDLE:        c_int = 10;
   // run this hook after every other hook which is defined
   pub const APR_HOOK_LAST:          c_int = 20;
   // run this hook last, after EVERYTHING
   pub const APR_HOOK_REALLY_LAST:   c_int = 30;

   extern "C" {
      pub fn apr_table_get(t: *const apr_table_t, key: *const c_char) -> *const c_char;
      pub fn apr_table_set(t: *mut apr_table_t, key: *const c_char, val: *const c_char) -> ();
      pub fn apr_table_elts(t: *const apr_table_t) -> *const apr_array_header_t;
      pub fn apr_pstrmemdup(p: *mut apr_pool_t, s: *const c_char, n: apr_size_t) -> *mut c_char;
      pub fn apr_version_string() -> *const c_char;
      pub fn apu_version_string() -> *const c_char;
   }

   pub fn dup_c_str<T: Into<Vec<u8>>>(pool: *mut apr_pool_t, data: T) -> *mut c_char {
      let bytes = data.into();

      unsafe {
         apr_pstrmemdup(
            pool,
            bytes.as_ptr() as *const c_char,
            bytes.len() as apr_size_t
         )
      }
   }

   #[repr(C)]
   pub struct apr_bucket_alloc_t;

   #[repr(C)]
   pub struct apr_bucket_brigade;

   #[repr(C)]
   pub struct apr_finfo_t;

   #[repr(C)]
   pub struct apr_pool_t;

   #[repr(C)]
   pub struct apr_sockaddr_t;

   #[repr(C)]
   pub struct apr_table_t;

   #[repr(C)]
   pub struct apr_thread_mutex_t;

   #[repr(C)]
   pub struct apr_thread_t;

   #[repr(C)]
   pub struct apr_uri_t;

   pub type apr_byte_t = c_uchar;
   pub type apr_int16_t = c_short;
   pub type apr_uint16_t = c_ushort;
   pub type apr_int32_t = c_int;
   pub type apr_uint32_t = c_uint;
   pub type apr_int64_t = c_long;
   pub type apr_uint64_t = c_ulong;
   pub type apr_size_t = c_ulong;
   pub type apr_ssize_t = c_long;
   pub type apr_off_t = c_long;
   pub type apr_socklen_t = c_uint;
   pub type apr_ino_t = c_ulong;
   pub type apr_uintptr_t = apr_uint64_t;
   pub type apr_status_t = c_int;
   pub type apr_signum_t = c_int;
   pub type apr_time_t = apr_int64_t;
   pub type apr_port_t = apr_uint16_t;

   #[repr(C)]
   pub struct apr_array_header_t {
      pub pool: *mut apr_pool_t,
      pub elt_size: c_int,
      pub nelts: c_int,
      pub nalloc: c_int,
      pub elts: *mut c_char,
   }

   #[repr(C)]
   pub struct apr_table_entry_t {
      pub key: *mut c_char,
      pub val: *mut c_char,
      pub key_checksum: apr_uint32_t,
   }
}

use std::ffi::CString;

use wrapper::{Wrapper, c_str_value};

pub type AprTable<'a> = Wrapper<'a, raw::apr_table_t>;

impl<'a> AprTable<'a> {
   pub fn get(&self, key: &'static str) -> Option<&'a str> {
      c_str_value(
         unsafe { raw::apr_table_get(self.raw, CString::new(key).unwrap().as_ptr()) }
      )
   }

   pub fn set(&mut self, key: &'static str, val: &'static str) {
      unsafe {
         raw::apr_table_set(
            self.raw,
            CString::new(key).unwrap().as_ptr(),
            CString::new(val).unwrap().as_ptr()
         )
      };
   }

   pub fn iter(&self) -> AprTableIter {
      let ptr = unsafe { raw::apr_table_elts(self.raw) };
      let raw: &raw::apr_array_header_t = unsafe { &*ptr };

      AprTableIter {
         array_header: raw,
         next_idx: 0
      }
   }
}

pub struct AprTableIter<'a> {
   array_header: &'a raw::apr_array_header_t,
   next_idx: usize,
}

impl<'a> Iterator for AprTableIter<'a> {
   type Item = (&'a str, &'a str);

   fn next(&mut self) -> Option<(&'a str, &'a str)> {
      if self.next_idx != self.array_header.nelts as usize {
         let mut elts = self.array_header.elts as *const raw::apr_table_entry_t;

         elts = unsafe { elts.offset(self.next_idx as isize) };
         self.next_idx += 1;

         let key = c_str_value(unsafe { (*elts).key }).unwrap();
         let val = c_str_value(unsafe { (*elts).val }).unwrap();

         Some((key, val))
      } else {
         None
      }
   }

   fn size_hint(&self) -> (usize, Option<usize>) {
      let rem = self.array_header.nelts as usize - self.next_idx;
      (rem, Some(rem))
   }
}

pub fn apr_version_string<'a>() -> Option<&'a str> {
   c_str_value(
      unsafe { raw::apr_version_string() }
   )
}

pub fn apu_version_string<'a>() -> Option<&'a str> {
   c_str_value(
      unsafe { raw::apu_version_string() }
   )
}
<|endoftext|>


<commit_before>#![allow(unsafe_code)]
/// We use this because we never use the weak count on the std
/// `Arc`, but we use a LOT of `Arc`'s, so the extra 8 bytes
/// turn into a huge overhead.
use std::{
    alloc::{alloc, dealloc, Layout},
    convert::TryFrom,
    fmt::{self, Debug},
    mem,
    ops::Deref,
    ptr,
    sync::atomic::{AtomicUsize, Ordering},
};

struct ArcInner<T: ?Sized> {
    rc: AtomicUsize,
    inner: T,
}

pub struct Arc<T: ?Sized> {
    ptr: *mut ArcInner<T>,
}

unsafe impl<T: Send + Sync + ?Sized> Send for Arc<T> {}
unsafe impl<T: Send + Sync + ?Sized> Sync for Arc<T> {}

impl<T: Debug> Debug for Arc<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        f.debug_map().entry(&"inner", unsafe { &(*self.ptr).inner }).finish()
    }
}

impl<T> Arc<T> {
    pub fn new(inner: T) -> Arc<T> {
        let bx = Box::new(ArcInner { inner, rc: AtomicUsize::new(1) });
        let ptr = Box::into_raw(bx);
        Arc { ptr }
    }

    pub fn copy_from_slice(s: &[T]) -> Arc<[T]> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());
        let data_width = mem::size_of::<T>().checked_add(s.len()).unwrap();

        let size = rc_width.checked_add(data_width).unwrap();

        let layout = Layout::from_size_align(size, align).unwrap();

        unsafe {
            let ptr = alloc(layout);

            assert!(!ptr.is_null(), "failed to allocate Arc");
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(ptr as _, AtomicUsize::new(1));

            let data_ptr = ptr.add(rc_width);
            ptr::copy_nonoverlapping(s.as_ptr(), data_ptr as _, s.len());

            let fat_ptr: *const ArcInner<[T]> = Arc::fatten(ptr, s.len());

            Arc { ptr: fat_ptr as *mut _ }
        }
    }

    /// <https://users.rust-lang.org/t/construct-fat-pointer-to-struct/29198/9>
    #[allow(trivial_casts)]
    fn fatten(data: *const u8, len: usize) -> *const ArcInner<[T]> {
        // Requirements of slice::from_raw_parts.
        assert!(!data.is_null());
        assert!(isize::try_from(len).is_ok());

        let slice =
            unsafe { core::slice::from_raw_parts(data as *const (), len) };
        slice as *const [()] as *const _
    }

    pub fn into_raw(arc: Arc<T>) -> *const T {
        let ptr = unsafe { &(*arc.ptr).inner };
        #[allow(clippy::mem_forget)]
        mem::forget(arc);
        ptr
    }

    pub unsafe fn from_raw(ptr: *const T) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let sub_ptr = (ptr as *const u8).sub(rc_width) as *mut ArcInner<T>;

        Arc { ptr: sub_ptr }
    }
}

impl<T: ?Sized> Arc<T> {
    pub fn strong_count(arc: &Arc<T>) -> usize {
        unsafe { (*arc.ptr).rc.load(Ordering::Acquire) }
    }

    pub fn get_mut(arc: &mut Arc<T>) -> Option<&mut T> {
        if Arc::strong_count(arc) == 1 {
            Some(unsafe { &mut arc.ptr.as_mut().unwrap().inner })
        } else {
            None
        }
    }
}

impl<T: ?Sized + Clone> Arc<T> {
    pub fn make_mut(arc: &mut Arc<T>) -> &mut T {
        if Arc::strong_count(arc) != 1 {
            *arc = Arc::new((**arc).clone());
            assert_eq!(Arc::strong_count(arc), 1);
        }
        Arc::get_mut(arc).unwrap()
    }
}

impl<T: Default> Default for Arc<T> {
    fn default() -> Arc<T> {
        Arc::new(T::default())
    }
}

impl<T: ?Sized> Clone for Arc<T> {
    fn clone(&self) -> Arc<T> {
        Arc { ptr: self.ptr }
    }
}

impl<T: ?Sized> Drop for Arc<T> {
    fn drop(&mut self) {
        unsafe {
            let rc = (*self.ptr).rc.fetch_sub(1, Ordering::Release) - 1;
            if rc == 0 {
                ptr::drop_in_place(self.ptr);
            }
        }
    }
}

impl<T: Copy> From<&[T]> for Arc<[T]> {
    #[inline]
    fn from(s: &[T]) -> Arc<[T]> {
        Arc::copy_from_slice(s)
    }
}

#[allow(clippy::fallible_impl_from)]
impl<T> From<Box<T>> for Arc<T> {
    #[inline]
    fn from(b: Box<T>) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let size = rc_width.checked_add(mem::size_of::<T>()).unwrap();

        unsafe {
            let dst_layout = Layout::from_size_align(size, align).unwrap();
            let dst = alloc(dst_layout);
            assert!(!dst.is_null(), "failed to allocate Arc");

            let data_ptr = dst.add(rc_width);
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(dst as _, AtomicUsize::new(1));
            let src = Box::into_raw(b);
            ptr::copy(src, data_ptr as *mut T, 1);

            // free the old box memory without running Drop
            let src_layout = Layout::new::<T>();
            dealloc(src as *mut u8, src_layout);

            Arc { ptr: dst as *mut _ }
        }
    }
}

impl<T> From<Vec<T>> for Arc<[T]> {
    #[inline]
    fn from(mut v: Vec<T>) -> Arc<[T]> {
        unsafe {
            let arc = Arc::copy_from_slice(&v);

            // Allow the Vec to free its memory, but not destroy its contents
            v.set_len(0);

            arc
        }
    }
}

impl<T: ?Sized> Deref for Arc<T> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &(*self.ptr).inner }
    }
}

impl<T: ?Sized> std::borrow::Borrow<T> for Arc<T> {
    fn borrow(&self) -> &T {
        &**self
    }
}

impl<T: ?Sized> AsRef<T> for Arc<T> {
    fn as_ref(&self) -> &T {
        &**self
    }
}
<commit_msg>Address divergentdave's feedback<commit_after>#![allow(unsafe_code)]
/// We use this because we never use the weak count on the std
/// `Arc`, but we use a LOT of `Arc`'s, so the extra 8 bytes
/// turn into a huge overhead.
use std::{
    alloc::{alloc, dealloc, Layout},
    convert::TryFrom,
    fmt::{self, Debug},
    mem,
    ops::Deref,
    ptr,
    sync::atomic::{AtomicUsize, Ordering},
};

struct ArcInner<T: ?Sized> {
    rc: AtomicUsize,
    inner: T,
}

pub struct Arc<T: ?Sized> {
    ptr: *mut ArcInner<T>,
}

unsafe impl<T: Send + Sync + ?Sized> Send for Arc<T> {}
unsafe impl<T: Send + Sync + ?Sized> Sync for Arc<T> {}

impl<T: Debug> Debug for Arc<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        f.debug_map().entry(&"inner", unsafe { &(*self.ptr).inner }).finish()
    }
}

impl<T> Arc<T> {
    pub fn new(inner: T) -> Arc<T> {
        let bx = Box::new(ArcInner { inner, rc: AtomicUsize::new(1) });
        let ptr = Box::into_raw(bx);
        Arc { ptr }
    }

    pub fn copy_from_slice(s: &[T]) -> Arc<[T]> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());
        let data_width = mem::size_of::<T>().checked_mul(s.len()).unwrap();

        let size = rc_width.checked_add(data_width).unwrap();

        let layout = Layout::from_size_align(size, align).unwrap();

        unsafe {
            let ptr = alloc(layout);

            assert!(!ptr.is_null(), "failed to allocate Arc");
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(ptr as _, AtomicUsize::new(1));

            let data_ptr = ptr.add(rc_width);
            ptr::copy_nonoverlapping(s.as_ptr(), data_ptr as _, s.len());

            let fat_ptr: *const ArcInner<[T]> = Arc::fatten(ptr, s.len());

            Arc { ptr: fat_ptr as *mut _ }
        }
    }

    /// <https://users.rust-lang.org/t/construct-fat-pointer-to-struct/29198/9>
    #[allow(trivial_casts)]
    fn fatten(data: *const u8, len: usize) -> *const ArcInner<[T]> {
        // Requirements of slice::from_raw_parts.
        assert!(!data.is_null());
        assert!(isize::try_from(len).is_ok());

        let slice =
            unsafe { core::slice::from_raw_parts(data as *const (), len) };
        slice as *const [()] as *const _
    }

    pub fn into_raw(arc: Arc<T>) -> *const T {
        let ptr = unsafe { &(*arc.ptr).inner };
        #[allow(clippy::mem_forget)]
        mem::forget(arc);
        ptr
    }

    pub unsafe fn from_raw(ptr: *const T) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let sub_ptr = (ptr as *const u8).sub(rc_width) as *mut ArcInner<T>;

        Arc { ptr: sub_ptr }
    }
}

impl<T: ?Sized> Arc<T> {
    pub fn strong_count(arc: &Arc<T>) -> usize {
        unsafe { (*arc.ptr).rc.load(Ordering::Acquire) }
    }

    pub fn get_mut(arc: &mut Arc<T>) -> Option<&mut T> {
        if Arc::strong_count(arc) == 1 {
            Some(unsafe { &mut arc.ptr.as_mut().unwrap().inner })
        } else {
            None
        }
    }
}

impl<T: ?Sized + Clone> Arc<T> {
    pub fn make_mut(arc: &mut Arc<T>) -> &mut T {
        if Arc::strong_count(arc) != 1 {
            *arc = Arc::new((**arc).clone());
            assert_eq!(Arc::strong_count(arc), 1);
        }
        Arc::get_mut(arc).unwrap()
    }
}

impl<T: Default> Default for Arc<T> {
    fn default() -> Arc<T> {
        Arc::new(T::default())
    }
}

impl<T: ?Sized> Clone for Arc<T> {
    fn clone(&self) -> Arc<T> {
        unsafe {
            (*self.ptr).rc.fetch_add(1, Ordering::Release);
        }
        Arc { ptr: self.ptr }
    }
}

impl<T: ?Sized> Drop for Arc<T> {
    fn drop(&mut self) {
        unsafe {
            let rc = (*self.ptr).rc.fetch_sub(1, Ordering::Release) - 1;
            if rc == 0 {
                ptr::drop_in_place(self.ptr);
            }
        }
    }
}

impl<T: Copy> From<&[T]> for Arc<[T]> {
    #[inline]
    fn from(s: &[T]) -> Arc<[T]> {
        Arc::copy_from_slice(s)
    }
}

#[allow(clippy::fallible_impl_from)]
impl<T> From<Box<T>> for Arc<T> {
    #[inline]
    fn from(b: Box<T>) -> Arc<T> {
        let align =
            std::cmp::max(mem::align_of::<T>(), mem::align_of::<AtomicUsize>());

        let rc_width = std::cmp::max(align, mem::size_of::<AtomicUsize>());

        let size = rc_width.checked_add(mem::size_of::<T>()).unwrap();

        unsafe {
            let dst_layout = Layout::from_size_align(size, align).unwrap();
            let dst = alloc(dst_layout);
            assert!(!dst.is_null(), "failed to allocate Arc");

            let data_ptr = dst.add(rc_width);
            #[allow(clippy::cast_ptr_alignment)]
            ptr::write(dst as _, AtomicUsize::new(1));
            let src = Box::into_raw(b);
            ptr::copy(src, data_ptr as *mut T, 1);

            // free the old box memory without running Drop
            let src_layout = Layout::new::<T>();
            dealloc(src as *mut u8, src_layout);

            Arc { ptr: dst as *mut _ }
        }
    }
}

impl<T> From<Vec<T>> for Arc<[T]> {
    #[inline]
    fn from(mut v: Vec<T>) -> Arc<[T]> {
        unsafe {
            let arc = Arc::copy_from_slice(&v);

            // Allow the Vec to free its memory, but not destroy its contents
            v.set_len(0);

            arc
        }
    }
}

impl<T: ?Sized> Deref for Arc<T> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &(*self.ptr).inner }
    }
}

impl<T: ?Sized> std::borrow::Borrow<T> for Arc<T> {
    fn borrow(&self) -> &T {
        &**self
    }
}

impl<T: ?Sized> AsRef<T> for Arc<T> {
    fn as_ref(&self) -> &T {
        &**self
    }
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate dent;

use clap::{App, Arg};
use dent::plot;
use dent::summary::Summary;
use dent::t_test::{SigLevel, TTest, welch_t_test};

use std::fs::File;
use std::path::Path;
use std::io::{self, BufRead, BufReader};

fn print_summary(s: &Summary) {
    println!("N\tMin\tMax\tMedian\tMean\tStdDev\tStdErr");
    println!(
        "{}\t{:0.2}\t{:0.2}\t{:0.2}\t{:0.2}\t{:0.2}\t{:0.2}",
        s.size(),
        s.min(),
        s.max(),
        s.median(),
        s.mean(),
        s.standard_deviation(),
        s.standard_error(),
    );
}

fn print_t_test(t_test: &TTest) {
    println!("T\tDF\tAlpha\tCrit\tRejectNull");
    println!(
        "{:0.3}\t{}\t{:0.3}\t{:0.3}\t{}",
        t_test.t,
        t_test.df,
        t_test.alpha,
        t_test.crit,
        t_test.reject,
    );
}

fn read_file(path: &str) -> Summary {
    let p = Path::new(path);
    let f = File::open(p).ok().unwrap();
    let r = BufReader::new(f);

    let data: Vec<f64> = r
        .lines()
        .map(|l| l.ok().unwrap().parse().unwrap())
        .collect();

    Summary::new(&data).unwrap()
}

fn parse_alpha(arg: &str) -> SigLevel {
    match arg {
        ".001" => SigLevel::Alpha001,
        ".005" => SigLevel::Alpha005,
        ".01"  => SigLevel::Alpha010,
        ".025" => SigLevel::Alpha025,
        ".05"  => SigLevel::Alpha050,
        ".1"   => SigLevel::Alpha100,
        _ => panic!(),
    }
}

fn summarize_stdin(draw_plot: bool) {
    let stdin = io::stdin();

    let data: Vec<f64> = stdin
        .lock()
        .lines()
        .map(|l| l.ok().unwrap().parse().unwrap())
        .collect();

    let s = Summary::new(&data).unwrap();

    if draw_plot {
        println!("{}\n", plot::summary_plot(&s, 80));
    }

    print_summary(&s);
}

fn t_test_files(file1: &str, file2: &str, alpha: SigLevel, draw_plot: bool) {
    let s1 = read_file(file1);
    let s2 = read_file(file2);

    let t_test = welch_t_test(&s1, &s2, alpha);

    if draw_plot {
        println!("{}\n", plot::comparison_plot(&s1, &s2, 80));
    }

    print_summary(&s1);
    println!();
    print_summary(&s2);
    println!();
    print_t_test(&t_test);
}

fn main() {
    let matches = App::new("dent")
        .version("0.1.0")
        .author("Joe Ranweiler <joe@lemma.co>")
        .about("A tiny tool for t-tests &c.")
        .arg(Arg::with_name("stdin")
             .short("s")
             .long("stdin")
             .help("Read and summarize data from stdin."))
        .arg(Arg::with_name("file1")
             .index(1)
             .value_name("FILE1")
             .takes_value(true)
             .required_unless("stdin")
             .help("1st file of sample data."))
        .arg(Arg::with_name("file2")
             .index(2)
             .value_name("FILE2")
             .takes_value(true)
             .required_unless("stdin")
             .help("2nd file of sample data."))
        .arg(Arg::with_name("alpha")
             .short("a")
             .long("alpha")
             .value_name("ALPHA")
             .help("Significance level α")
             .takes_value(true)
             .default_value(".05"))
        .arg(Arg::with_name("plot")
             .short("p")
             .long("plot")
             .help("Draw a box-and-whisker plot."))
        .get_matches();

    let draw_plot = matches.is_present("plot");
    let use_stdin = matches.is_present("stdin");

    if use_stdin {
        summarize_stdin(draw_plot);
    } else {
        let alpha = parse_alpha(matches.value_of("alpha").unwrap());
        let file1 = matches.value_of("file1").unwrap();
        let file2 = matches.value_of("file2").unwrap();

        t_test_files(file1, file2, alpha, draw_plot);
    }
}
<commit_msg>Add flag for ASCII-only plotting<commit_after>extern crate clap;
extern crate dent;

use clap::{App, Arg};
use dent::plot;
use dent::summary::Summary;
use dent::t_test::{SigLevel, TTest, welch_t_test};

use std::fs::File;
use std::path::Path;
use std::io::{self, BufRead, BufReader};

fn print_summary(s: &Summary) {
    println!("N\tMin\tMax\tMedian\tMean\tStdDev\tStdErr");
    println!(
        "{}\t{:0.2}\t{:0.2}\t{:0.2}\t{:0.2}\t{:0.2}\t{:0.2}",
        s.size(),
        s.min(),
        s.max(),
        s.median(),
        s.mean(),
        s.standard_deviation(),
        s.standard_error(),
    );
}

fn print_t_test(t_test: &TTest) {
    println!("T\tDF\tAlpha\tCrit\tRejectNull");
    println!(
        "{:0.3}\t{}\t{:0.3}\t{:0.3}\t{}",
        t_test.t,
        t_test.df,
        t_test.alpha,
        t_test.crit,
        t_test.reject,
    );
}

fn read_file(path: &str) -> Summary {
    let p = Path::new(path);
    let f = File::open(p).ok().unwrap();
    let r = BufReader::new(f);

    let data: Vec<f64> = r
        .lines()
        .map(|l| l.ok().unwrap().parse().unwrap())
        .collect();

    Summary::new(&data).unwrap()
}

fn parse_alpha(arg: &str) -> SigLevel {
    match arg {
        ".001" => SigLevel::Alpha001,
        ".005" => SigLevel::Alpha005,
        ".01"  => SigLevel::Alpha010,
        ".025" => SigLevel::Alpha025,
        ".05"  => SigLevel::Alpha050,
        ".1"   => SigLevel::Alpha100,
        _ => panic!(),
    }
}

fn summarize_stdin(draw_plot: bool, ascii: bool) {
    let stdin = io::stdin();

    let data: Vec<f64> = stdin
        .lock()
        .lines()
        .map(|l| l.ok().unwrap().parse().unwrap())
        .collect();

    let s = Summary::new(&data).unwrap();

    if draw_plot {
        if ascii {
            println!("{}\n", plot::ascii_summary_plot(&s, 80));
        } else {
            println!("{}\n", plot::summary_plot(&s, 80));
        }
    }

    print_summary(&s);
}

fn t_test_files(
    file1: &str,
    file2: &str,
    alpha: SigLevel,
    draw_plot: bool,
    ascii: bool,
) {
    let s1 = read_file(file1);
    let s2 = read_file(file2);

    let t_test = welch_t_test(&s1, &s2, alpha);

    if draw_plot {
        println!("{}\n", plot::comparison_plot(&s1, &s2, 80, ascii));
    }

    print_summary(&s1);
    println!();
    print_summary(&s2);
    println!();
    print_t_test(&t_test);
}

fn main() {
    let matches = App::new("dent")
        .version("0.1.0")
        .author("Joe Ranweiler <joe@lemma.co>")
        .about("A tiny tool for t-tests &c.")
        .arg(Arg::with_name("stdin")
             .short("s")
             .long("stdin")
             .help("Read and summarize data from stdin."))
        .arg(Arg::with_name("file1")
             .index(1)
             .value_name("FILE1")
             .takes_value(true)
             .required_unless("stdin")
             .help("1st file of sample data."))
        .arg(Arg::with_name("file2")
             .index(2)
             .value_name("FILE2")
             .takes_value(true)
             .required_unless("stdin")
             .help("2nd file of sample data."))
        .arg(Arg::with_name("alpha")
             .short("a")
             .long("alpha")
             .value_name("ALPHA")
             .help("Significance level α")
             .takes_value(true)
             .default_value(".05"))
        .arg(Arg::with_name("plot")
             .short("p")
             .long("plot")
             .help("Draw a box-and-whisker plot."))
        .arg(Arg::with_name("ascii")
             .long("ascii")
             .help("Use only ASCII characters in plots."))
        .get_matches();

    let ascii = matches.is_present("ascii");
    let draw_plot = matches.is_present("plot");
    let use_stdin = matches.is_present("stdin");

    if use_stdin {
        summarize_stdin(draw_plot, ascii);
    } else {
        let alpha = parse_alpha(matches.value_of("alpha").unwrap());
        let file1 = matches.value_of("file1").unwrap();
        let file2 = matches.value_of("file2").unwrap();

        t_test_files(file1, file2, alpha, draw_plot, ascii);
    }
}
<|endoftext|>