<commit_before>extern crate bytecount;
extern crate glob;
extern crate handlebars;
#[macro_use]
extern crate lazy_static;
extern crate lazycell;
#[macro_use]
extern crate log;
extern crate num_cpus;
extern crate rand;
extern crate regex;
extern crate scoped_threadpool;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;
extern crate simple_logger;
extern crate syntect;
extern crate time;
extern crate toml;
extern crate typed_arena;
extern crate walkdir;

mod directives;
mod evaluator;
mod highlighter;
mod init;
mod lex;
mod page;
mod parse;
mod theme;
mod toctree;

use std::collections::HashMap;
use std::convert::From;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::ops::DerefMut;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::{env, mem, process};
use evaluator::{Evaluator, Worker};
use page::{Page, Slug};
use toctree::TocTree;
use directives::logic;
use scoped_threadpool::Pool;

#[derive(Debug)]
enum LinkError {
    UndefinedReference,
    TemplateError(handlebars::RenderError),
    IOError(io::Error),
}

impl From<handlebars::RenderError> for LinkError {
    fn from(orig: handlebars::RenderError) -> Self {
        LinkError::TemplateError(orig)
    }
}

impl From<io::Error> for LinkError {
    fn from(orig: io::Error) -> Self {
        LinkError::IOError(orig)
    }
}

#[derive(Deserialize)]
struct RawConfig {
    syntax_theme: Option<String>,
    theme: Option<PathBuf>,
    content_dir: Option<PathBuf>,
    output: Option<PathBuf>,
    templates: HashMap<String, String>,
    theme_constants: Option<serde_json::map::Map<String, serde_json::Value>>,
}

struct Project {
    verbose: bool,
    theme: PathBuf,
    content_dir: PathBuf,
    output: PathBuf,
    templates: Vec<(glob::Pattern, String)>,
    theme_constants: serde_json::map::Map<String, serde_json::Value>,
    syntax_theme: String,

    pretty_url: bool,
}

impl Project {
    fn read_toml(path: &Path) -> Result<Project, ()> {
        let mut file = File::open(path).or(Err(()))?;
        let mut data = String::new();
        file.read_to_string(&mut data).or(Err(()))?;
        let config: RawConfig = toml::from_str(&data).or(Err(()))?;

        let theme = config.theme.ok_or(())?;

        let path_patterns: Result<Vec<_>, ()> = config
            .templates
            .iter()
            .map(|(k, v)| {
                let pattern = match glob::Pattern::new(k) {
                    Ok(p) => p,
                    Err(_) => return Err(()),
                };
                Ok((pattern, v.to_owned()))
            })
            .collect();

        let path_patterns = path_patterns.or(Err(()))?;

        let syntax_theme = config
            .syntax_theme
            .unwrap_or_else(|| highlighter::DEFAULT_SYNTAX_THEME.to_owned());

        Ok(Project {
            verbose: false,
            theme,
            content_dir: config
                .content_dir
                .unwrap_or_else(|| PathBuf::from("content")),
            output: config.output.unwrap_or_else(|| PathBuf::from("build")),
            templates: path_patterns,
            theme_constants: config
                .theme_constants
                .unwrap_or_else(serde_json::map::Map::new),
            syntax_theme,
            pretty_url: true,
        })
    }

    fn build_file(&self, worker: &mut Worker, path: &Path) -> Result<Page, ()> {
        debug!("Compiling {}", worker.get_slug());

        let node = match worker.parser.parse(path) {
            Ok(n) => n,
            Err(msg) => {
                error!("Failed to parse '{}': {}", path.to_string_lossy(), msg);
                return Err(());
            }
        };

        let output = worker.evaluate(&node);

        let page = Page {
            source_path: path.to_owned(),
            slug: worker.get_slug().clone(),
            body: output,
            theme_config: worker.theme_config.clone(),
        };

        Ok(page)
    }

    fn link_file(
        &self,
        evaluator: &Evaluator,
        page: &Page,
        renderer: &theme::Renderer,
    ) -> Result<(), LinkError> {
        debug!("Linking {}", &page.slug);

        // Find the template that matches this path
        let template_name = self.templates
            .iter()
            .find(|&&(ref pat, _)| pat.matches_path(&page.source_path))
            .map(|&(_, ref name)| name.as_ref())
            .unwrap_or("default");

        let new_body = match evaluator.substitute(page) {
            Ok(s) => s,
            Err(_) => {
                return Err(LinkError::UndefinedReference);
            }
        };

        let rendered = renderer.render(template_name, &self.theme_constants, page, &new_body)?;
        let output_path = page.slug.create_output_path(&self.output, self.pretty_url);
        let output_dir = output_path.parent().expect("Couldn't get output directory");

        fs::create_dir_all(output_dir)?;
        let mut file = File::create(&output_path)?;
        file.write_all(rendered.as_bytes())?;

        Ok(())
    }
}

fn build_project(project: Project, evaluator: Evaluator) {
    let num_cpus = num_cpus::get();
    let project = Arc::new(project);
    let evaluator = Arc::new(evaluator);
    let titles: Arc<Mutex<HashMap<Slug, String>>> = Arc::new(Mutex::new(HashMap::new()));
    let pending_pages: Arc<Mutex<Vec<Page>>> = Arc::new(Mutex::new(vec![]));

    debug!("Crawling source directory");

    let mut paths = vec![];
    for entry in walkdir::WalkDir::new(&project.content_dir) {
        let entry = entry.expect("Failed to walk dir");
        if !entry.file_type().is_file() {
            continue;
        }

        if entry.path().extension() != Some("rocket".as_ref()) {
            continue;
        }

        paths.push(entry.path().to_owned());
    }

    debug!("Compiling with {} workers", num_cpus);
    let paths = Arc::new(paths);
    let chunk_size = (paths.len() as f32 / num_cpus as f32).ceil() as usize;
    let chunks: Vec<_> = paths.chunks(chunk_size).map(|x| x.to_owned()).collect();
    let mut threads = Vec::with_capacity(chunks.len());
    for chunk in chunks {
        let project = Arc::clone(&project);
        let evaluator = Arc::clone(&evaluator);
        let titles = Arc::clone(&titles);
        let pending_pages = Arc::clone(&pending_pages);

        let thread = std::thread::spawn(move || {
            let mut worker = Worker::new_with_options(&evaluator, &project.syntax_theme);

            for path in chunk {
                let slug = path.strip_prefix(&project.content_dir)
                    .expect("Failed to get output path");
                let dir = slug.parent().unwrap();
                let stem = slug.file_stem().unwrap();
                let slug = Slug::new(dir.join(stem).to_string_lossy().as_ref().to_owned());
                worker.set_slug(slug);

                match project.build_file(&mut worker, &path) {
                    Ok(page) => {
                        titles
                            .lock()
                            .unwrap()
                            .insert(page.slug.to_owned(), page.title());
                        pending_pages.lock().unwrap().push(page);
                    }
                    Err(_) => {
                        error!("Failed to build {}", path.to_string_lossy());
                    }
                }
            }
        });

        threads.push(thread);
    }

    for thread in threads {
        thread.join().unwrap();
    }

    let mut toctree = {
        let mut txn = evaluator.toctree.write().unwrap();
        mem::replace(txn.deref_mut(), TocTree::new_empty())
    };

    toctree.finish(titles.lock().unwrap().deref());

    let theme = theme::Theme::load(&project.theme).expect("Failed to load theme");

    let renderer = Arc::new(
        theme::Renderer::new(theme, Arc::new(toctree)).expect("Failed to construct renderer"),
    );

    debug!("Linking with {} workers", num_cpus);

    let mut pool = Pool::new(num_cpus as u32);
    pool.scoped(move |scoped| {
        let mut pending_pages = pending_pages.lock().unwrap();
        for page in pending_pages.drain(0..) {
            let project = Arc::clone(&project);
            let evaluator = Arc::clone(&evaluator);
            let renderer = Arc::clone(&renderer);

            scoped.execute(move || {
                project
                    .link_file(&evaluator, &page, &renderer)
                    .expect("Failed to link page");
            });
        }
    });
}

fn build(verbose: bool) {
    let mut config =
        Project::read_toml(Path::new("config.toml")).expect("Failed to open config.toml");

    config.verbose = verbose;

    let mut evaluator = Evaluator::new();
    evaluator.register_prelude("code", Box::new(directives::Code));
    evaluator.register_prelude("table", Box::new(directives::Dummy));
    evaluator.register_prelude("version", Box::new(directives::Version::new("3.4.0")));
    evaluator.register_prelude(
        "note",
        Box::new(directives::Admonition::new("Note", "note")),
    );
    evaluator.register_prelude(
        "warning",
        Box::new(directives::Admonition::new("Warning", "warning")),
    );
    evaluator.register_prelude("define-template", Box::new(directives::DefineTemplate));
    evaluator.register_prelude("definition-list", Box::new(directives::DefinitionList));
    evaluator.register_prelude("steps", Box::new(directives::Steps));
    evaluator.register_prelude("concat", Box::new(directives::Concat));
    evaluator.register_prelude("include", Box::new(directives::Include));
    evaluator.register_prelude("import", Box::new(directives::Import));
    evaluator.register_prelude("null", Box::new(directives::Dummy));
    evaluator.register_prelude("let", Box::new(directives::Let));
    evaluator.register_prelude("define", Box::new(directives::Define));
    evaluator.register_prelude("theme-config", Box::new(directives::ThemeConfig));
    evaluator.register_prelude("toctree", Box::new(directives::TocTree));
    evaluator.register_prelude("define-ref", Box::new(directives::RefDefDirective));
    evaluator.register_prelude("ref", Box::new(directives::RefDirective));
    evaluator.register_prelude("link", Box::new(directives::Link));
    evaluator.register_prelude("figure", Box::new(directives::Figure));
    evaluator.register_prelude("ul", Box::new(directives::List::new("ul")));
    evaluator.register_prelude("ol", Box::new(directives::List::new("ol")));

    // Formatting
    evaluator.register_prelude("``", Box::new(directives::FormattingMarker::new("code")));
    evaluator.register_prelude("**", Box::new(directives::FormattingMarker::new("strong")));
    evaluator.register_prelude("__", Box::new(directives::FormattingMarker::new("em")));

    // Headers
    evaluator.register_prelude("h1", Box::new(directives::Heading::new(1)));
    evaluator.register_prelude("h2", Box::new(directives::Heading::new(2)));
    evaluator.register_prelude("h3", Box::new(directives::Heading::new(3)));
    evaluator.register_prelude("h4", Box::new(directives::Heading::new(4)));
    evaluator.register_prelude("h5", Box::new(directives::Heading::new(5)));
    evaluator.register_prelude("h6", Box::new(directives::Heading::new(6)));

    // Logic operations
    evaluator.register_prelude("if", Box::new(logic::If));
    evaluator.register_prelude("not", Box::new(logic::Not));
    evaluator.register_prelude("=", Box::new(logic::Equals));
    evaluator.register_prelude("!=", Box::new(logic::NotEquals));

    let start_time = time::precise_time_ns();
    build_project(config, evaluator);

    info!(
        "Took {} seconds",
        (time::precise_time_ns() - start_time) as f64 / (f64::from(1_000_000_000))
    );
}

const DESCRIPTION_BUILD: &'static str =
    "Build the Rocket project in the current working directory.";
const DESCRIPTION_NEW: &'static str = "Create an empty Rocket project.";
const HELP_VERBOSE: &'static str = "Increase logging verbosity.";

enum ArgMode {
    Root,
    New,
    Build,
}

fn main() {
    let args = env::args().skip(1);
    let mut verbose = false;
    let mut new_name: Option<String> = None;
    let mut mode = ArgMode::Root;

    let help = |code| -> ! {
        println!("Usage:\n  rocket [-h, OPTS...] {{ new | build }} ...\n");
        println!("Description:\n  The Rocket documentation build system.\n");
        println!(
            "Subcommands:\n  new\n    {}\n  build\n    {}\n",
            DESCRIPTION_NEW,
            DESCRIPTION_BUILD
        );
        println!("Optional arguments:");
        println!("  --help, -h\n    Print this message and exit.\n");
        println!("  --version, -V\n    Print version string and exit.\n");

        process::exit(code);
    };

    let help_build = |code| -> ! {
        println!("Usage:\n  rocket build [-h, OPTS...]\n");
        println!("Description:\n  {}\n", DESCRIPTION_BUILD);
        println!("Optional arguments:");
        println!("  --verbose, -v\n    {}\n", HELP_VERBOSE);
        println!("  --help, -h\n    Print this message and exit.\n");

        process::exit(code);
    };

    let help_new = |code| -> ! {
        println!("Usage:\n  rocket new [-h, OPTS...] name\n");
        println!("Description:\n  {}\n", DESCRIPTION_NEW);
        println!("Positional arguments:\n  name\n    The name of the project to create.\n");
        println!("Optional arguments:");
        println!("  --verbose, -v\n    {}\n", HELP_VERBOSE);
        println!("  --help, -h\n    Print this message and exit.\n");

        process::exit(code);
    };

    for arg in args {
        match mode {
            ArgMode::Root => match arg.as_ref() {
                "-h" | "--help" => help(0),
                "-V" | "--version" => {
                    println!(
                        "{}",
                        option_env!("CARGO_PKG_VERSION").unwrap_or("<unknown>")
                    );
                    return;
                }
                "-v" | "--verbose" => verbose = true,
                "build" => mode = ArgMode::Build,
                "new" => mode = ArgMode::New,
                _ => help(1),
            },
            ArgMode::New => {
                let alphanumeric = arg.chars().all(|c| c.is_alphanumeric());
                match arg.as_ref() {
                    "-h" | "--help" => help_new(0),
                    "-v" | "--verbose" => verbose = true,
                    n if alphanumeric => new_name = Some(n.to_owned()),
                    _ => help_new(1),
                }
            }
            ArgMode::Build => match arg.as_ref() {
                "-h" | "--help" => help_build(0),
                "-v" | "--verbose" => verbose = true,
                _ => help_build(1),
            },
        }
    }

    let loglevel = if verbose {
        log::LogLevel::Debug
    } else {
        log::LogLevel::Info
    };

    simple_logger::init_with_level(loglevel).expect("Failed to initialize logger");

    match mode {
        ArgMode::Root => help(1),
        ArgMode::New => init::init(&new_name.unwrap_or_else(|| help_new(1))),
        ArgMode::Build => build(verbose),
    }
}
<commit_msg>Be more specific when unwrapping thread.join() result<commit_after>extern crate bytecount;
extern crate glob;
extern crate handlebars;
#[macro_use]
extern crate lazy_static;
extern crate lazycell;
#[macro_use]
extern crate log;
extern crate num_cpus;
extern crate rand;
extern crate regex;
extern crate scoped_threadpool;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;
extern crate simple_logger;
extern crate syntect;
extern crate time;
extern crate toml;
extern crate typed_arena;
extern crate walkdir;

mod directives;
mod evaluator;
mod highlighter;
mod init;
mod lex;
mod page;
mod parse;
mod theme;
mod toctree;

use std::collections::HashMap;
use std::convert::From;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::ops::DerefMut;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::{env, mem, process};
use evaluator::{Evaluator, Worker};
use page::{Page, Slug};
use toctree::TocTree;
use directives::logic;
use scoped_threadpool::Pool;

#[derive(Debug)]
enum LinkError {
    UndefinedReference,
    TemplateError(handlebars::RenderError),
    IOError(io::Error),
}

impl From<handlebars::RenderError> for LinkError {
    fn from(orig: handlebars::RenderError) -> Self {
        LinkError::TemplateError(orig)
    }
}

impl From<io::Error> for LinkError {
    fn from(orig: io::Error) -> Self {
        LinkError::IOError(orig)
    }
}

#[derive(Deserialize)]
struct RawConfig {
    syntax_theme: Option<String>,
    theme: Option<PathBuf>,
    content_dir: Option<PathBuf>,
    output: Option<PathBuf>,
    templates: HashMap<String, String>,
    theme_constants: Option<serde_json::map::Map<String, serde_json::Value>>,
}

struct Project {
    verbose: bool,
    theme: PathBuf,
    content_dir: PathBuf,
    output: PathBuf,
    templates: Vec<(glob::Pattern, String)>,
    theme_constants: serde_json::map::Map<String, serde_json::Value>,
    syntax_theme: String,

    pretty_url: bool,
}

impl Project {
    fn read_toml(path: &Path) -> Result<Project, ()> {
        let mut file = File::open(path).or(Err(()))?;
        let mut data = String::new();
        file.read_to_string(&mut data).or(Err(()))?;
        let config: RawConfig = toml::from_str(&data).or(Err(()))?;

        let theme = config.theme.ok_or(())?;

        let path_patterns: Result<Vec<_>, ()> = config
            .templates
            .iter()
            .map(|(k, v)| {
                let pattern = match glob::Pattern::new(k) {
                    Ok(p) => p,
                    Err(_) => return Err(()),
                };
                Ok((pattern, v.to_owned()))
            })
            .collect();

        let path_patterns = path_patterns.or(Err(()))?;

        let syntax_theme = config
            .syntax_theme
            .unwrap_or_else(|| highlighter::DEFAULT_SYNTAX_THEME.to_owned());

        Ok(Project {
            verbose: false,
            theme,
            content_dir: config
                .content_dir
                .unwrap_or_else(|| PathBuf::from("content")),
            output: config.output.unwrap_or_else(|| PathBuf::from("build")),
            templates: path_patterns,
            theme_constants: config
                .theme_constants
                .unwrap_or_else(serde_json::map::Map::new),
            syntax_theme,
            pretty_url: true,
        })
    }

    fn build_file(&self, worker: &mut Worker, path: &Path) -> Result<Page, ()> {
        debug!("Compiling {}", worker.get_slug());

        let node = match worker.parser.parse(path) {
            Ok(n) => n,
            Err(msg) => {
                error!("Failed to parse '{}': {}", path.to_string_lossy(), msg);
                return Err(());
            }
        };

        let output = worker.evaluate(&node);

        let page = Page {
            source_path: path.to_owned(),
            slug: worker.get_slug().clone(),
            body: output,
            theme_config: worker.theme_config.clone(),
        };

        Ok(page)
    }

    fn link_file(
        &self,
        evaluator: &Evaluator,
        page: &Page,
        renderer: &theme::Renderer,
    ) -> Result<(), LinkError> {
        debug!("Linking {}", &page.slug);

        // Find the template that matches this path
        let template_name = self.templates
            .iter()
            .find(|&&(ref pat, _)| pat.matches_path(&page.source_path))
            .map(|&(_, ref name)| name.as_ref())
            .unwrap_or("default");

        let new_body = match evaluator.substitute(page) {
            Ok(s) => s,
            Err(_) => {
                return Err(LinkError::UndefinedReference);
            }
        };

        let rendered = renderer.render(template_name, &self.theme_constants, page, &new_body)?;
        let output_path = page.slug.create_output_path(&self.output, self.pretty_url);
        let output_dir = output_path.parent().expect("Couldn't get output directory");

        fs::create_dir_all(output_dir)?;
        let mut file = File::create(&output_path)?;
        file.write_all(rendered.as_bytes())?;

        Ok(())
    }
}

fn build_project(project: Project, evaluator: Evaluator) {
    let num_cpus = num_cpus::get();
    let project = Arc::new(project);
    let evaluator = Arc::new(evaluator);
    let titles: Arc<Mutex<HashMap<Slug, String>>> = Arc::new(Mutex::new(HashMap::new()));
    let pending_pages: Arc<Mutex<Vec<Page>>> = Arc::new(Mutex::new(vec![]));

    debug!("Crawling source directory");

    let mut paths = vec![];
    for entry in walkdir::WalkDir::new(&project.content_dir) {
        let entry = entry.expect("Failed to walk dir");
        if !entry.file_type().is_file() {
            continue;
        }

        if entry.path().extension() != Some("rocket".as_ref()) {
            continue;
        }

        paths.push(entry.path().to_owned());
    }

    debug!("Compiling with {} workers", num_cpus);
    let paths = Arc::new(paths);
    let chunk_size = (paths.len() as f32 / num_cpus as f32).ceil() as usize;
    let chunks: Vec<_> = paths.chunks(chunk_size).map(|x| x.to_owned()).collect();
    let mut threads = Vec::with_capacity(chunks.len());
    for chunk in chunks {
        let project = Arc::clone(&project);
        let evaluator = Arc::clone(&evaluator);
        let titles = Arc::clone(&titles);
        let pending_pages = Arc::clone(&pending_pages);

        let thread = std::thread::spawn(move || {
            let mut worker = Worker::new_with_options(&evaluator, &project.syntax_theme);

            for path in chunk {
                let slug = path.strip_prefix(&project.content_dir)
                    .expect("Failed to get output path");
                let dir = slug.parent().unwrap();
                let stem = slug.file_stem().unwrap();
                let slug = Slug::new(dir.join(stem).to_string_lossy().as_ref().to_owned());
                worker.set_slug(slug);

                match project.build_file(&mut worker, &path) {
                    Ok(page) => {
                        titles
                            .lock()
                            .unwrap()
                            .insert(page.slug.to_owned(), page.title());
                        pending_pages.lock().unwrap().push(page);
                    }
                    Err(_) => {
                        error!("Failed to build {}", path.to_string_lossy());
                    }
                }
            }
        });

        threads.push(thread);
    }

    for thread in threads {
        thread.join().expect("At least one compilation worker panicked");
    }

    let mut toctree = {
        let mut txn = evaluator.toctree.write().unwrap();
        mem::replace(txn.deref_mut(), TocTree::new_empty())
    };

    toctree.finish(titles.lock().unwrap().deref());

    let theme = theme::Theme::load(&project.theme).expect("Failed to load theme");

    let renderer = Arc::new(
        theme::Renderer::new(theme, Arc::new(toctree)).expect("Failed to construct renderer"),
    );

    debug!("Linking with {} workers", num_cpus);

    let mut pool = Pool::new(num_cpus as u32);
    pool.scoped(move |scoped| {
        let mut pending_pages = pending_pages.lock().unwrap();
        for page in pending_pages.drain(0..) {
            let project = Arc::clone(&project);
            let evaluator = Arc::clone(&evaluator);
            let renderer = Arc::clone(&renderer);

            scoped.execute(move || {
                project
                    .link_file(&evaluator, &page, &renderer)
                    .expect("Failed to link page");
            });
        }
    });
}

fn build(verbose: bool) {
    let mut config =
        Project::read_toml(Path::new("config.toml")).expect("Failed to open config.toml");

    config.verbose = verbose;

    let mut evaluator = Evaluator::new();
    evaluator.register_prelude("code", Box::new(directives::Code));
    evaluator.register_prelude("table", Box::new(directives::Dummy));
    evaluator.register_prelude("version", Box::new(directives::Version::new("3.4.0")));
    evaluator.register_prelude(
        "note",
        Box::new(directives::Admonition::new("Note", "note")),
    );
    evaluator.register_prelude(
        "warning",
        Box::new(directives::Admonition::new("Warning", "warning")),
    );
    evaluator.register_prelude("define-template", Box::new(directives::DefineTemplate));
    evaluator.register_prelude("definition-list", Box::new(directives::DefinitionList));
    evaluator.register_prelude("steps", Box::new(directives::Steps));
    evaluator.register_prelude("concat", Box::new(directives::Concat));
    evaluator.register_prelude("include", Box::new(directives::Include));
    evaluator.register_prelude("import", Box::new(directives::Import));
    evaluator.register_prelude("null", Box::new(directives::Dummy));
    evaluator.register_prelude("let", Box::new(directives::Let));
    evaluator.register_prelude("define", Box::new(directives::Define));
    evaluator.register_prelude("theme-config", Box::new(directives::ThemeConfig));
    evaluator.register_prelude("toctree", Box::new(directives::TocTree));
    evaluator.register_prelude("define-ref", Box::new(directives::RefDefDirective));
    evaluator.register_prelude("ref", Box::new(directives::RefDirective));
    evaluator.register_prelude("link", Box::new(directives::Link));
    evaluator.register_prelude("figure", Box::new(directives::Figure));
    evaluator.register_prelude("ul", Box::new(directives::List::new("ul")));
    evaluator.register_prelude("ol", Box::new(directives::List::new("ol")));

    // Formatting
    evaluator.register_prelude("``", Box::new(directives::FormattingMarker::new("code")));
    evaluator.register_prelude("**", Box::new(directives::FormattingMarker::new("strong")));
    evaluator.register_prelude("__", Box::new(directives::FormattingMarker::new("em")));

    // Headers
    evaluator.register_prelude("h1", Box::new(directives::Heading::new(1)));
    evaluator.register_prelude("h2", Box::new(directives::Heading::new(2)));
    evaluator.register_prelude("h3", Box::new(directives::Heading::new(3)));
    evaluator.register_prelude("h4", Box::new(directives::Heading::new(4)));
    evaluator.register_prelude("h5", Box::new(directives::Heading::new(5)));
    evaluator.register_prelude("h6", Box::new(directives::Heading::new(6)));

    // Logic operations
    evaluator.register_prelude("if", Box::new(logic::If));
    evaluator.register_prelude("not", Box::new(logic::Not));
    evaluator.register_prelude("=", Box::new(logic::Equals));
    evaluator.register_prelude("!=", Box::new(logic::NotEquals));

    let start_time = time::precise_time_ns();
    build_project(config, evaluator);

    info!(
        "Took {} seconds",
        (time::precise_time_ns() - start_time) as f64 / (f64::from(1_000_000_000))
    );
}

const DESCRIPTION_BUILD: &'static str =
    "Build the Rocket project in the current working directory.";
const DESCRIPTION_NEW: &'static str = "Create an empty Rocket project.";
const HELP_VERBOSE: &'static str = "Increase logging verbosity.";

enum ArgMode {
    Root,
    New,
    Build,
}

fn main() {
    let args = env::args().skip(1);
    let mut verbose = false;
    let mut new_name: Option<String> = None;
    let mut mode = ArgMode::Root;

    let help = |code| -> ! {
        println!("Usage:\n  rocket [-h, OPTS...] {{ new | build }} ...\n");
        println!("Description:\n  The Rocket documentation build system.\n");
        println!(
            "Subcommands:\n  new\n    {}\n  build\n    {}\n",
            DESCRIPTION_NEW,
            DESCRIPTION_BUILD
        );
        println!("Optional arguments:");
        println!("  --help, -h\n    Print this message and exit.\n");
        println!("  --version, -V\n    Print version string and exit.\n");

        process::exit(code);
    };

    let help_build = |code| -> ! {
        println!("Usage:\n  rocket build [-h, OPTS...]\n");
        println!("Description:\n  {}\n", DESCRIPTION_BUILD);
        println!("Optional arguments:");
        println!("  --verbose, -v\n    {}\n", HELP_VERBOSE);
        println!("  --help, -h\n    Print this message and exit.\n");

        process::exit(code);
    };

    let help_new = |code| -> ! {
        println!("Usage:\n  rocket new [-h, OPTS...] name\n");
        println!("Description:\n  {}\n", DESCRIPTION_NEW);
        println!("Positional arguments:\n  name\n    The name of the project to create.\n");
        println!("Optional arguments:");
        println!("  --verbose, -v\n    {}\n", HELP_VERBOSE);
        println!("  --help, -h\n    Print this message and exit.\n");

        process::exit(code);
    };

    for arg in args {
        match mode {
            ArgMode::Root => match arg.as_ref() {
                "-h" | "--help" => help(0),
                "-V" | "--version" => {
                    println!(
                        "{}",
                        option_env!("CARGO_PKG_VERSION").unwrap_or("<unknown>")
                    );
                    return;
                }
                "-v" | "--verbose" => verbose = true,
                "build" => mode = ArgMode::Build,
                "new" => mode = ArgMode::New,
                _ => help(1),
            },
            ArgMode::New => {
                let alphanumeric = arg.chars().all(|c| c.is_alphanumeric());
                match arg.as_ref() {
                    "-h" | "--help" => help_new(0),
                    "-v" | "--verbose" => verbose = true,
                    n if alphanumeric => new_name = Some(n.to_owned()),
                    _ => help_new(1),
                }
            }
            ArgMode::Build => match arg.as_ref() {
                "-h" | "--help" => help_build(0),
                "-v" | "--verbose" => verbose = true,
                _ => help_build(1),
            },
        }
    }

    let loglevel = if verbose {
        log::LogLevel::Debug
    } else {
        log::LogLevel::Info
    };

    simple_logger::init_with_level(loglevel).expect("Failed to initialize logger");

    match mode {
        ArgMode::Root => help(1),
        ArgMode::New => init::init(&new_name.unwrap_or_else(|| help_new(1))),
        ArgMode::Build => build(verbose),
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::env;
use std::fmt;
use std::io;
use std::io::Write;
use std::sync::Mutex;

extern crate syn;
use syn::*;

mod rust_lib;

#[derive(Debug)]
struct ConvertedType {
    /// The type converted to C (e.g. `uint32_t`)
    prefix: String,
    /// Stuff that might need to go after the identifier (e.g. `[3]` for an array).
    postfix: String,
}

impl ConvertedType {
    fn new<T: Into<String>>(pre: T, post: String) -> Self {
        ConvertedType {
            prefix: pre.into(),
            postfix: post
        }
    }

    fn format_with_ident(&self, ident: &Ident) -> String {
        format!("{} {}{}", &self.prefix, &ident.to_string(), &self.postfix)
    }
}

impl Into<String> for ConvertedType {
    fn into(self) -> String {
        let mut str = String::from(self.prefix);
        str.push_str(&self.postfix);
        str
    }
}

impl From<String> for ConvertedType {
    fn from(str: String) -> ConvertedType {
        ConvertedType {
            prefix: str,
            postfix: String::new()
        }
    }
}

impl fmt::Display for ConvertedType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}{}", &self.prefix, &self.postfix)
    }
}

fn has_attribute(target: MetaItem, attrs: &Vec<Attribute>) -> bool {
    return attrs
               .iter()
               .any(|ref attr| attr.style == AttrStyle::Outer && attr.value == target);
}

fn has_no_mangle(attrs: &Vec<Attribute>) -> bool {
    has_attribute(MetaItem::Word(Ident::new("no_mangle")), attrs)
}

fn wr_func_body(attrs: &Vec<Attribute>) -> String {
    if has_attribute(MetaItem::Word(Ident::new("destructor_safe")), attrs) {
        String::from("WR_DESTRUCTOR_SAFE_FUNC")
    } else {
        String::from("WR_FUNC")
    }
}

fn is_repr_c(attrs: &Vec<Attribute>) -> bool {
    let repr_args = vec![NestedMetaItem::MetaItem(MetaItem::Word(Ident::new("C")))];
    has_attribute(MetaItem::List(Ident::new("repr"), repr_args), attrs)
}

fn is_repr_u32(attrs: &Vec<Attribute>) -> bool {
    let repr_args = vec![NestedMetaItem::MetaItem(MetaItem::Word(Ident::new("u32")))];
    has_attribute(MetaItem::List(Ident::new("repr"), repr_args), attrs)
}

fn is_c_abi(abi: &Option<Abi>) -> bool {
    abi == &Some(Abi::Named(String::from("C")))
}

fn map_path(p: &Path) -> String {
    let l = p.segments[0].ident.to_string();
    let mut c = match l.as_ref() {
        "usize" => "size_t".to_string(),
        "u8" => "uint8_t".to_string(),
        "u32" => "uint32_t".to_string(),
        "u64" => "uint64_t".to_string(),
        "i8" => "int8_t".to_string(),
        "i32" => "int32_t".to_string(),
        "i64" => "int64_t".to_string(),
        "f32" => "float".to_string(),
        "c_void" => "void".to_string(),
        _ => l,
    };
    if let PathParameters::AngleBracketed(ref d) = p.segments[0].parameters {
        let template_args = d.types
            .iter()
            .map(|ty| map_ty(ty).into())
            .collect::<Vec<String>>()
            .join(", ");
        if !template_args.is_empty() {
            c.push('<');
            c.push_str(&template_args);
            c.push('>');
        }
    }
    c
}

fn map_mut_ty(mut_ty: &MutTy) -> ConvertedType {
    map_ty(&mut_ty.ty)
}

fn map_ty(ty: &Ty) -> ConvertedType {
    match ty {
        &Ty::Path(_, ref p) => ConvertedType::from(map_path(p)),
        &Ty::Ptr(ref p) => ConvertedType::from(format!("{}*", map_ty(&p.ty))),
        &Ty::Rptr(_, ref mut_ty) => ConvertedType::from(format!("{}*", map_mut_ty(mut_ty))),
        &Ty::Array(ref p, ConstExpr::Lit(Lit::Int(sz, _))) => ConvertedType::new(map_ty(&p), format!("[{}]", sz)),
        _ => ConvertedType::from(format!("unknown {:?}", ty)),
    }
}

fn map_return_type(ret: &FunctionRetTy) -> String {
    match ret {
        &FunctionRetTy::Default => "void".to_string(),
        &FunctionRetTy::Ty(ref ty) => map_ty(ty).into(),
    }
}

fn map_arg(f: &FnArg) -> String {
    match f {
        &FnArg::Captured(Pat::Ident(_, ref ident, _), ref ty) => map_ty(ty).format_with_ident(ident),
        _ => "unknown".to_string(),
    }
}

fn map_field(f: &Field) -> String {
    let mut ret = String::from("  ");
    let converted = map_ty(&f.ty);
    ret.push_str(&converted.format_with_ident(f.ident.as_ref().expect("Struct fields must have idents")));
    ret.push_str(";\n");
    ret
}

fn map_generic_param(t: &TyParam) -> String {
    let mut ret = String::from("typename ");
    ret.push_str(&t.ident.to_string());
    ret
}

fn fold_enum_variants(accum: (String, i32), v: &Variant) -> (String, i32) {
    // `accum` contains the combined string of converted enum variants so far, to which
    // we will append the converted version of `v`. The other thing in `accum` is the
    // value of the previous variant. If `v` has an explicit value we keep that, otherwise
    // we increment the value of the previous variant to get the new one. This is all
    // so that we properly support enums with explicitly-specified and discontinuous
    // values.
    let mut ret = accum.0;
    ret.push_str("  ");
    ret.push_str(&v.ident.to_string());
    ret.push_str(" = ");
    let new_value = match &v.discriminant {
        &None => accum.1 + 1,
        &Some(ConstExpr::Lit(Lit::Int(ref specified_value, _))) => *specified_value as i32,
        &Some(_) => {
            // we don't handle this yet, so just put in something that will fail C compilation
            writeln!(io::stderr(), "warning, unsupported enum discriminant").unwrap();
            ret.push_str("???");
            accum.1 + 1
        }
    };
    ret.push_str(&new_value.to_string());
    ret.push_str(",\n");
    (ret, new_value)
}

/// A structure to store all the results of converting Rust stuff to C
/// stuff, so that we can pick and choose what we output.
struct ConversionResults {
    /// This holds the C function signatures of no_mangle rust functions,
    /// in the order that they were encountered in Rust.
    funcs: Vec<String>,
    /// This holds the C conversions of repr(C) structs and repr(u32) enums
    /// from Rust. The 'ds' in the name stands for 'data structures'. The
    /// key of the map is the name of the type, and the value is the C code.
    c_ds: HashMap<String, String>,
    /// This holds the dependency tree. If a struct contains another struct
    /// or an enum it will be in this map. The key is the dependent type,
    /// the list of values are all the things it depends on.
    _dep_tree: HashMap<String, Vec<String>>,
}

fn main() {
    let p = env::args().nth(1).unwrap();

    let results = Mutex::new(ConversionResults {
        funcs: Vec::new(),
        c_ds: HashMap::new(),
        _dep_tree: HashMap::new(),
    });

    rust_lib::parse(p, &|mod_name, items| {
        for item in items {
            match item.node {
                ItemKind::Fn(ref decl,
                             ref _unsafe,
                             ref _const,
                             ref abi,
                             ref _generic,
                             ref _block) => {
                    writeln!(io::stderr(), "processing function {}::{}", mod_name, &item.ident).unwrap();
                    if has_no_mangle(&item.attrs) && is_c_abi(&abi) {
                        results.lock().unwrap().funcs.push(
                            format!("WR_INLINE {}\n{}({})\n{};\n",
                                    map_return_type(&decl.output),
                                    item.ident,
                                    decl.inputs
                                        .iter()
                                        .map(map_arg)
                                        .collect::<Vec<_>>()
                                        .join(", "),
                                    wr_func_body(&item.attrs)));
                    }
                }
                ItemKind::Struct(ref variant,
                                 ref generics) => {
                    writeln!(io::stderr(), "processing struct {}::{}", mod_name, &item.ident).unwrap();
                    if is_repr_c(&item.attrs) {
                        if let &VariantData::Struct(ref fields) = variant {
                            let mut c_struct = String::new();
                            if !generics.ty_params.is_empty() {
                                c_struct.push_str(&(
                                    format!("template<{}>",
                                            generics.ty_params
                                                    .iter()
                                                    .map(map_generic_param)
                                                    .collect::<Vec<_>>()
                                                    .join(", "))));
                            }
                            c_struct.push_str(&(
                                format!("struct {} {{\n{}}};\n",
                                        item.ident,
                                        fields.iter()
                                              .map(map_field)
                                              .collect::<String>())));
                            results.lock().unwrap().c_ds.insert(
                                item.ident.to_string(),
                                c_struct);
                        }
                    }
                }
                ItemKind::Enum(ref variants, ref _generics) => {
                    writeln!(io::stderr(), "processing enum {}::{}", mod_name, &item.ident).unwrap();
                    if is_repr_u32(&item.attrs) {
                        let c_enum = format!("enum class {}: uint32_t {{\n{}\n  Sentinel /* this must be last for serialization purposes. */\n}};\n",
                                             item.ident,
                                             variants.iter()
                                                     .fold((String::new(), -1), fold_enum_variants)
                                                     .0);
                        results.lock().unwrap().c_ds.insert(
                            item.ident.to_string(),
                            c_enum);
                    }
                }
                _ => {}
            }
        }
    });
    for (_, c_stuff) in &results.lock().unwrap().c_ds {
        println!("{}", c_stuff);
    }
    for func in &results.lock().unwrap().funcs {
        println!("{}", func);
    }
}
<commit_msg>Make map_return_type return a ConvertedType<commit_after>use std::collections::HashMap;
use std::env;
use std::fmt;
use std::io;
use std::io::Write;
use std::sync::Mutex;

extern crate syn;
use syn::*;

mod rust_lib;

#[derive(Debug)]
struct ConvertedType {
    /// The type converted to C (e.g. `uint32_t`)
    prefix: String,
    /// Stuff that might need to go after the identifier (e.g. `[3]` for an array).
    postfix: String,
}

impl ConvertedType {
    fn new<T: Into<String>>(pre: T, post: String) -> Self {
        ConvertedType {
            prefix: pre.into(),
            postfix: post
        }
    }

    fn format_with_ident(&self, ident: &Ident) -> String {
        format!("{} {}{}", &self.prefix, &ident.to_string(), &self.postfix)
    }
}

impl Into<String> for ConvertedType {
    fn into(self) -> String {
        let mut str = String::from(self.prefix);
        str.push_str(&self.postfix);
        str
    }
}

impl From<String> for ConvertedType {
    fn from(str: String) -> ConvertedType {
        ConvertedType {
            prefix: str,
            postfix: String::new()
        }
    }
}

impl fmt::Display for ConvertedType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}{}", &self.prefix, &self.postfix)
    }
}

fn has_attribute(target: MetaItem, attrs: &Vec<Attribute>) -> bool {
    return attrs
               .iter()
               .any(|ref attr| attr.style == AttrStyle::Outer && attr.value == target);
}

fn has_no_mangle(attrs: &Vec<Attribute>) -> bool {
    has_attribute(MetaItem::Word(Ident::new("no_mangle")), attrs)
}

fn wr_func_body(attrs: &Vec<Attribute>) -> String {
    if has_attribute(MetaItem::Word(Ident::new("destructor_safe")), attrs) {
        String::from("WR_DESTRUCTOR_SAFE_FUNC")
    } else {
        String::from("WR_FUNC")
    }
}

fn is_repr_c(attrs: &Vec<Attribute>) -> bool {
    let repr_args = vec![NestedMetaItem::MetaItem(MetaItem::Word(Ident::new("C")))];
    has_attribute(MetaItem::List(Ident::new("repr"), repr_args), attrs)
}

fn is_repr_u32(attrs: &Vec<Attribute>) -> bool {
    let repr_args = vec![NestedMetaItem::MetaItem(MetaItem::Word(Ident::new("u32")))];
    has_attribute(MetaItem::List(Ident::new("repr"), repr_args), attrs)
}

fn is_c_abi(abi: &Option<Abi>) -> bool {
    abi == &Some(Abi::Named(String::from("C")))
}

fn map_path(p: &Path) -> String {
    let l = p.segments[0].ident.to_string();
    let mut c = match l.as_ref() {
        "usize" => "size_t".to_string(),
        "u8" => "uint8_t".to_string(),
        "u32" => "uint32_t".to_string(),
        "u64" => "uint64_t".to_string(),
        "i8" => "int8_t".to_string(),
        "i32" => "int32_t".to_string(),
        "i64" => "int64_t".to_string(),
        "f32" => "float".to_string(),
        "c_void" => "void".to_string(),
        _ => l,
    };
    if let PathParameters::AngleBracketed(ref d) = p.segments[0].parameters {
        let template_args = d.types
            .iter()
            .map(|ty| map_ty(ty).into())
            .collect::<Vec<String>>()
            .join(", ");
        if !template_args.is_empty() {
            c.push('<');
            c.push_str(&template_args);
            c.push('>');
        }
    }
    c
}

fn map_mut_ty(mut_ty: &MutTy) -> ConvertedType {
    map_ty(&mut_ty.ty)
}

fn map_ty(ty: &Ty) -> ConvertedType {
    match ty {
        &Ty::Path(_, ref p) => ConvertedType::from(map_path(p)),
        &Ty::Ptr(ref p) => ConvertedType::from(format!("{}*", map_ty(&p.ty))),
        &Ty::Rptr(_, ref mut_ty) => ConvertedType::from(format!("{}*", map_mut_ty(mut_ty))),
        &Ty::Array(ref p, ConstExpr::Lit(Lit::Int(sz, _))) => ConvertedType::new(map_ty(&p), format!("[{}]", sz)),
        _ => ConvertedType::from(format!("unknown {:?}", ty)),
    }
}

fn map_return_type(ret: &FunctionRetTy) -> ConvertedType {
    match ret {
        &FunctionRetTy::Default => ConvertedType::from("void".to_string()),
        &FunctionRetTy::Ty(ref ty) => map_ty(ty),
    }
}

fn map_arg(f: &FnArg) -> String {
    match f {
        &FnArg::Captured(Pat::Ident(_, ref ident, _), ref ty) => map_ty(ty).format_with_ident(ident),
        _ => "unknown".to_string(),
    }
}

fn map_field(f: &Field) -> String {
    let mut ret = String::from("  ");
    let converted = map_ty(&f.ty);
    ret.push_str(&converted.format_with_ident(f.ident.as_ref().expect("Struct fields must have idents")));
    ret.push_str(";\n");
    ret
}

fn map_generic_param(t: &TyParam) -> String {
    let mut ret = String::from("typename ");
    ret.push_str(&t.ident.to_string());
    ret
}

fn fold_enum_variants(accum: (String, i32), v: &Variant) -> (String, i32) {
    // `accum` contains the combined string of converted enum variants so far, to which
    // we will append the converted version of `v`. The other thing in `accum` is the
    // value of the previous variant. If `v` has an explicit value we keep that, otherwise
    // we increment the value of the previous variant to get the new one. This is all
    // so that we properly support enums with explicitly-specified and discontinuous
    // values.
    let mut ret = accum.0;
    ret.push_str("  ");
    ret.push_str(&v.ident.to_string());
    ret.push_str(" = ");
    let new_value = match &v.discriminant {
        &None => accum.1 + 1,
        &Some(ConstExpr::Lit(Lit::Int(ref specified_value, _))) => *specified_value as i32,
        &Some(_) => {
            // we don't handle this yet, so just put in something that will fail C compilation
            writeln!(io::stderr(), "warning, unsupported enum discriminant").unwrap();
            ret.push_str("???");
            accum.1 + 1
        }
    };
    ret.push_str(&new_value.to_string());
    ret.push_str(",\n");
    (ret, new_value)
}

/// A structure to store all the results of converting Rust stuff to C
/// stuff, so that we can pick and choose what we output.
struct ConversionResults {
    /// This holds the C function signatures of no_mangle rust functions,
    /// in the order that they were encountered in Rust.
    funcs: Vec<String>,
    /// This holds the C conversions of repr(C) structs and repr(u32) enums
    /// from Rust. The 'ds' in the name stands for 'data structures'. The
    /// key of the map is the name of the type, and the value is the C code.
    c_ds: HashMap<String, String>,
    /// This holds the dependency tree. If a struct contains another struct
    /// or an enum it will be in this map. The key is the dependent type,
    /// the list of values are all the things it depends on.
    _dep_tree: HashMap<String, Vec<String>>,
}

fn main() {
    let p = env::args().nth(1).unwrap();

    let results = Mutex::new(ConversionResults {
        funcs: Vec::new(),
        c_ds: HashMap::new(),
        _dep_tree: HashMap::new(),
    });

    rust_lib::parse(p, &|mod_name, items| {
        for item in items {
            match item.node {
                ItemKind::Fn(ref decl,
                             ref _unsafe,
                             ref _const,
                             ref abi,
                             ref _generic,
                             ref _block) => {
                    writeln!(io::stderr(), "processing function {}::{}", mod_name, &item.ident).unwrap();
                    if has_no_mangle(&item.attrs) && is_c_abi(&abi) {
                        results.lock().unwrap().funcs.push(
                            format!("WR_INLINE {}\n{}({})\n{};\n",
                                    map_return_type(&decl.output),
                                    item.ident,
                                    decl.inputs
                                        .iter()
                                        .map(map_arg)
                                        .collect::<Vec<_>>()
                                        .join(", "),
                                    wr_func_body(&item.attrs)));
                    }
                }
                ItemKind::Struct(ref variant,
                                 ref generics) => {
                    writeln!(io::stderr(), "processing struct {}::{}", mod_name, &item.ident).unwrap();
                    if is_repr_c(&item.attrs) {
                        if let &VariantData::Struct(ref fields) = variant {
                            let mut c_struct = String::new();
                            if !generics.ty_params.is_empty() {
                                c_struct.push_str(&(
                                    format!("template<{}>",
                                            generics.ty_params
                                                    .iter()
                                                    .map(map_generic_param)
                                                    .collect::<Vec<_>>()
                                                    .join(", "))));
                            }
                            c_struct.push_str(&(
                                format!("struct {} {{\n{}}};\n",
                                        item.ident,
                                        fields.iter()
                                              .map(map_field)
                                              .collect::<String>())));
                            results.lock().unwrap().c_ds.insert(
                                item.ident.to_string(),
                                c_struct);
                        }
                    }
                }
                ItemKind::Enum(ref variants, ref _generics) => {
                    writeln!(io::stderr(), "processing enum {}::{}", mod_name, &item.ident).unwrap();
                    if is_repr_u32(&item.attrs) {
                        let c_enum = format!("enum class {}: uint32_t {{\n{}\n  Sentinel /* this must be last for serialization purposes. */\n}};\n",
                                             item.ident,
                                             variants.iter()
                                                     .fold((String::new(), -1), fold_enum_variants)
                                                     .0);
                        results.lock().unwrap().c_ds.insert(
                            item.ident.to_string(),
                            c_enum);
                    }
                }
                _ => {}
            }
        }
    });
    for (_, c_stuff) in &results.lock().unwrap().c_ds {
        println!("{}", c_stuff);
    }
    for func in &results.lock().unwrap().funcs {
        println!("{}", func);
    }
}
<|endoftext|>


<commit_before>extern crate ar;
extern crate clap;
extern crate libflate;
extern crate tar;
extern crate tempfile;
extern crate time as crates_time;
extern crate xz2;
extern crate zip;

use std::fs;
use std::io;
use std::time;

use std::rc::Rc;
use std::vec::Vec;

use clap::{Arg, App};

use libflate::gzip;
use tempfile::tempfile;

// magic:
use std::io::Seek;
use std::io::Write;

#[derive(Debug)]
struct Node<T> {
    next: Option<Rc<Node<T>>>,
    value: T,
}

impl<T> Node<T> {
    fn head(obj: T) -> Rc<Node<T>> {
        Rc::new(Node {
            next: None,
            value: obj,
        })
    }

    fn plus(what: &Rc<Node<T>>, obj: T) -> Rc<Node<T>> {
        Rc::new(Node {
            next: Some(what.clone()),
            value: obj
        })
    }
}

struct OutputTo {
    path: Rc<Node<String>>,
    size: u64,
    atime: u64,
    mtime: u64,
    ctime: u64,
    btime: u64,
}

impl OutputTo {
    fn warn(&self, msg: String) {
        writeln!(io::stderr(), "TODO: {}", msg).expect("stderr");
    }

    fn raw<F: io::Read>(&self, mut file: &mut F) -> io::Result<()> {
        let stdout = io::stdout();
        let mut stdout = stdout.lock();
        writeln!(stdout, "{:?} {} {} {} {} {}",
            self.path,
            self.size,
            self.atime, self.mtime, self.ctime, self.btime
        )?;

        io::copy(&mut file, &mut stdout).and_then(move |written|
            if written != self.size {
                Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    format!("expecting to write {} but wrote {}", self.size, written)))
            } else {
                Ok(())
            })
    }
}

impl OutputTo {
    fn from_file(path: &str, fd: &fs::File) -> io::Result<OutputTo> {
        let meta = fd.metadata()?;
        Ok(OutputTo {
            path: Node::head(path.to_string()),
            size: meta.len(),
            atime: meta.accessed().map(simple_time_sys)?,
            mtime: meta.modified().map(simple_time_sys)?,
            ctime: simple_time_ctime(&meta),
            btime: simple_time_btime(&meta)?,
        })
    }

    fn with_path(&self, path: &str) -> OutputTo {
        OutputTo {
            path: Node::plus(&self.path, path.to_string()),
            size: 0,
            atime: 0,
            mtime: 0,
            ctime: 0,
            btime: 0,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum FileType {
    GZip,
    Zip,
    Tar,
    Ar,
    BZip2,
    Xz,
    Other,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Status {
    Done,
    Rewind,
}

fn read_octal(bytes: &[u8]) -> Option<u32> {
    let mut start = 0;
    while start < bytes.len() && b' ' == bytes[start] {
        start += 1;
    }

    let mut end = bytes.len() - 1;
    while end > start && (b' ' == bytes[end] || 0 == bytes[end]) {
        end -= 1;
    }

    if let Ok(string) = std::str::from_utf8(&bytes[start..(end+1)]) {
        if let Ok(val) = u32::from_str_radix(string, 8) {
            return Some(val);
        }
    }
    None
}

fn simple_time(dur: time::Duration) -> u64 {
    dur.as_secs().checked_mul(1_000_000_000)
        .map_or(0, |nanos| nanos + dur.subsec_nanos() as u64)
}

fn simple_time_sys(val: time::SystemTime) -> u64 {
    val.duration_since(time::UNIX_EPOCH).map(simple_time).unwrap_or(0)
}

fn simple_time_tm(val: crates_time::Tm) -> u64 {
    let timespec = val.to_timespec();
    simple_time(time::Duration::new(timespec.sec as u64, timespec.nsec as u32))
}

fn simple_time_btime(val: &fs::Metadata) -> io::Result<u64> {
    match val.created() {
        Ok(time) => Ok(simple_time_sys(time)),
        // "Other" is how "unsupported" is represented here; ew.
        Err(ref e) if e.kind() == io::ErrorKind::Other => Ok(0),
        Err(other) => Err(other),
    }
}

// TODO: I really feel this should be exposed by Rust, without cfg.
fn simple_time_ctime(val: &fs::Metadata) -> u64 {
    #[cfg(linux)] {
        use std::os::linux::fs::MetadataExt;
        let ctime: i64 = val.st_ctime();
        if ctime <= 0 {
            0
        } else {
            ctime as u64
        }
    }

    #[cfg(not(linux))] {
        0
    }
}

fn is_probably_tar(header: &[u8]) -> bool {
    if header.len() < 512 {
        return false;
    }

    if let Some(checksum) = read_octal(&header[148..156]) {
        let mut sum: u32 = (b' ' as u32) * 8;
        for i in 0..148 {
            sum += header[i] as u32;
        }
        for i in 156..512 {
            sum += header[i] as u32;
        }

        if checksum == sum {
            return true;
        }
    }

    return false;
}

fn identify<'a>(fd: &mut Box<io::BufRead + 'a>) -> io::Result<FileType> {
    let header = fd.fill_buf()?;
    if header.len() >= 20
        && 0x1f == header[0] && 0x8b == header[1] {
        Ok(FileType::GZip)
   } else if header.len() >= 152
        && b'P' == header[0] && b'K' == header[1]
        && 0x03 == header[2] && 0x04 == header[3] {
        Ok(FileType::Zip)
    } else if header.len() > 257 + 10
        && b'u' == header[257] && b's' == header[258]
        && b't' == header[259] && b'a' == header[260]
        && b'r' == header[261]
        && (
            (0 == header[262] && b'0' == header[263] && b'0' == header[264]) ||
            (b' ' == header[262] && b' ' == header[263] && 0 == header[264])
        ) {
        Ok(FileType::Tar)
    } else if header.len() > 8
        && b'!' == header[0] && b'<' == header[1]
        && b'a' == header[2] && b'r' == header[3]
        && b'c' == header[4] && b'h' == header[5]
        && b'>' == header[6] && b'\n' == header[7] {
        Ok(FileType::Ar)
    } else if header.len() > 40
        && b'B' == header[0] && b'Z' == header[1]
        && b'h' == header[2] && 0x31 == header[3]
        && 0x41 == header[4] && 0x59 == header[5]
        && 0x26 == header[6] {
        Ok(FileType::BZip2)
    } else if header.len() > 6
        && 0xfd == header[0] && b'7' == header[1]
        && b'z' == header[2] && b'X' == header[3]
        && b'Z' == header[4] && 0 == header[5] {
        Ok(FileType::Xz)
    } else if is_probably_tar(header) {
        Ok(FileType::Tar)
    } else {
        Ok(FileType::Other)
    }
}

struct BoxReader<'a> {
    fd: Box<io::BufRead + 'a>
}

impl<'a> io::Read for BoxReader<'a> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        self.fd.read(&mut buf)
    }
}

fn count_bytes<'a, R: io::Read>(mut fd: &mut R) -> io::Result<u64> {
    let mut block = [0u8; 4096];
    let mut count = 0u64;
    loop {
        let found = fd.read(&mut block[..])?;
        if 0 == found {
            break;
        }
        count += found as u64;
    }
    Ok(count)
}

fn unpack_or_raw<'a, F>(
    fd: Box<io::BufRead + 'a>,
    output: &OutputTo,
    fun: F) -> io::Result<Status>
where F: FnOnce(Box<io::BufRead + 'a>) -> io::Result<Box<io::BufRead + 'a>>
{
    match fun(fd) {
        Ok(inner) => unpack(inner, output),
        Err(e) => {
            output.warn(format!(
                    "thought we could unpack '{:?}' but we couldn't: {}",
                    output.path, e));
            Ok(Status::Rewind)
        }
    }
}

fn unpack<'a>(mut fd: Box<io::BufRead + 'a>, output: &OutputTo) -> io::Result<Status> {
    match identify(&mut fd)? {
        FileType::GZip => {
            unpack_or_raw(fd, output,
                |fd| gzip::Decoder::new(fd).map(
                    |dec| Box::new(io::BufReader::new(dec)) as Box<io::BufRead>))
        },
        FileType::Xz => {
            unpack_or_raw(fd, output,
                |fd| Ok(Box::new(io::BufReader::new(xz2::bufread::XzDecoder::new(fd)))))
        },
        FileType::Ar if output.path.value.ends_with(".deb") => {
            let mut decoder = ar::Archive::new(fd);
            while let Some(entry) = decoder.next_entry() {
                let entry = entry?;
                let new_output = output.with_path(entry.header().identifier());
                unpack(Box::new(io::BufReader::new(entry)), &new_output)?;
            }
            Ok(Status::Done)
        },
        FileType::Tar => {
            let mut decoder = tar::Archive::new(fd);
            for entry in decoder.entries()? {
                let entry = entry?;
                let new_output = output.with_path(entry.path()?.to_str().expect("valid utf-8"));
                unpack(Box::new(io::BufReader::new(entry)), &new_output)?;
            }
            Ok(Status::Done)
        },
        FileType::Zip => {
            let mut temp = tempfile()?;
            io::copy(&mut BoxReader { fd }, &mut temp)?;
            let mut zip = zip::ZipArchive::new(temp)?;

            for i in 0..zip.len() {

                // well, this has gone rather poorly
                let mut new_output = output.with_path(zip.by_index(i)?.name());

                let file_result = {
                    let entry = zip.by_index(i)?;

                    new_output.size = entry.size();
                    new_output.mtime = simple_time_tm(entry.last_modified());
                    let reader = Box::new(io::BufReader::new(entry));
                    unpack(reader, &new_output)?
                };
                match file_result {
                    Status::Done => (),
                    Status::Rewind => {
                        // TODO: update output? ? handling? ???
                        new_output.raw(&mut zip.by_index(i)?)?
                    },
                };
            }

            Ok(Status::Done)
        },

        // TODO: unimplemented!()
        _ => {
            Ok(Status::Rewind)
        },
    }
}

fn output_raw(file: &mut fs::File, output: &OutputTo) -> io::Result<()> {
    file.seek(io::SeekFrom::Start(0))?;
    output.raw(file)?;
    Ok(())
}

fn process_real_path(path: &str) -> io::Result<()> {
    let mut file = fs::File::open(path).unwrap();
    let output = OutputTo::from_file(
        path,
        &file
    )?;

    let attempt = {
        let read = io::BufReader::new(&file);
        unpack(Box::new(read), &output)
    };

    match attempt {
        Ok(Status::Done) => Ok(()),
        Ok(Status::Rewind) => output_raw(&mut file, &output),
        Err(e) => Err(e)
    }
}

fn real_main() -> u8 {

    let matches = App::new("contentin")
                    .arg(Arg::with_name("to-tar")
                         .long("to-tar")
                         .help("emit a tar file")
                         .required(true))
                    .arg(Arg::with_name("INPUT")
                         .required(true)
                         .multiple(true))
                    .get_matches();

    for path in matches.values_of("INPUT").unwrap() {
        if let Err(e) = process_real_path(path) {
            if let Err(_) = writeln!(io::stderr(), "fatal: processing '{}': {}", path, e) {
                return 6;
            }

            return 4;
        }
    }
    return 0;
}

fn main() {
    std::process::exit(real_main() as i32)
}
<commit_msg>vec up the path<commit_after>extern crate ar;
extern crate clap;
extern crate libflate;
extern crate tar;
extern crate tempfile;
extern crate time as crates_time;
extern crate xz2;
extern crate zip;

use std::fs;
use std::io;
use std::time;

use std::rc::Rc;
use std::vec::Vec;

use clap::{Arg, App};

use libflate::gzip;
use tempfile::tempfile;

// magic:
use std::io::Seek;
use std::io::Write;

#[derive(Debug)]
struct Node<T> {
    next: Option<Rc<Node<T>>>,
    value: T,
}

impl<T: Clone> Node<T> {
    fn head(obj: T) -> Rc<Node<T>> {
        Rc::new(Node {
            next: None,
            value: obj,
        })
    }

    fn plus(what: &Rc<Node<T>>, obj: T) -> Rc<Node<T>> {
        Rc::new(Node {
            next: Some(what.clone()),
            value: obj
        })
    }

    fn to_vec(&self) -> Vec<T> {
        let mut ret = Vec::new();
        let mut val = self;
        loop {
            ret.push(val.value.clone());
            if let Some(ref next) = val.next {
                val = next;
            } else {
                break;
            }
        }
        ret.reverse();
        ret
    }
}

struct OutputTo {
    path: Rc<Node<String>>,
    size: u64,
    atime: u64,
    mtime: u64,
    ctime: u64,
    btime: u64,
}

impl OutputTo {
    fn warn(&self, msg: String) {
        writeln!(io::stderr(), "TODO: {}", msg).expect("stderr");
    }

    fn raw<F: io::Read>(&self, mut file: &mut F) -> io::Result<()> {
        let stdout = io::stdout();
        let mut stdout = stdout.lock();
        writeln!(stdout, "{:?} {} {} {} {} {}",
            self.path.to_vec(),
            self.size,
            self.atime, self.mtime, self.ctime, self.btime
        )?;

        io::copy(&mut file, &mut stdout).and_then(move |written|
            if written != self.size {
                Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    format!("expecting to write {} but wrote {}", self.size, written)))
            } else {
                Ok(())
            })
    }
}

impl OutputTo {
    fn from_file(path: &str, fd: &fs::File) -> io::Result<OutputTo> {
        let meta = fd.metadata()?;
        Ok(OutputTo {
            path: Node::head(path.to_string()),
            size: meta.len(),
            atime: meta.accessed().map(simple_time_sys)?,
            mtime: meta.modified().map(simple_time_sys)?,
            ctime: simple_time_ctime(&meta),
            btime: simple_time_btime(&meta)?,
        })
    }

    fn with_path(&self, path: &str) -> OutputTo {
        OutputTo {
            path: Node::plus(&self.path, path.to_string()),
            size: 0,
            atime: 0,
            mtime: 0,
            ctime: 0,
            btime: 0,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum FileType {
    GZip,
    Zip,
    Tar,
    Ar,
    BZip2,
    Xz,
    Other,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Status {
    Done,
    Rewind,
}

fn read_octal(bytes: &[u8]) -> Option<u32> {
    let mut start = 0;
    while start < bytes.len() && b' ' == bytes[start] {
        start += 1;
    }

    let mut end = bytes.len() - 1;
    while end > start && (b' ' == bytes[end] || 0 == bytes[end]) {
        end -= 1;
    }

    if let Ok(string) = std::str::from_utf8(&bytes[start..(end+1)]) {
        if let Ok(val) = u32::from_str_radix(string, 8) {
            return Some(val);
        }
    }
    None
}

fn simple_time(dur: time::Duration) -> u64 {
    dur.as_secs().checked_mul(1_000_000_000)
        .map_or(0, |nanos| nanos + dur.subsec_nanos() as u64)
}

fn simple_time_sys(val: time::SystemTime) -> u64 {
    val.duration_since(time::UNIX_EPOCH).map(simple_time).unwrap_or(0)
}

fn simple_time_tm(val: crates_time::Tm) -> u64 {
    let timespec = val.to_timespec();
    simple_time(time::Duration::new(timespec.sec as u64, timespec.nsec as u32))
}

fn simple_time_btime(val: &fs::Metadata) -> io::Result<u64> {
    match val.created() {
        Ok(time) => Ok(simple_time_sys(time)),
        // "Other" is how "unsupported" is represented here; ew.
        Err(ref e) if e.kind() == io::ErrorKind::Other => Ok(0),
        Err(other) => Err(other),
    }
}

// TODO: I really feel this should be exposed by Rust, without cfg.
fn simple_time_ctime(val: &fs::Metadata) -> u64 {
    #[cfg(linux)] {
        use std::os::linux::fs::MetadataExt;
        let ctime: i64 = val.st_ctime();
        if ctime <= 0 {
            0
        } else {
            ctime as u64
        }
    }

    #[cfg(not(linux))] {
        0
    }
}

fn is_probably_tar(header: &[u8]) -> bool {
    if header.len() < 512 {
        return false;
    }

    if let Some(checksum) = read_octal(&header[148..156]) {
        let mut sum: u32 = (b' ' as u32) * 8;
        for i in 0..148 {
            sum += header[i] as u32;
        }
        for i in 156..512 {
            sum += header[i] as u32;
        }

        if checksum == sum {
            return true;
        }
    }

    return false;
}

fn identify<'a>(fd: &mut Box<io::BufRead + 'a>) -> io::Result<FileType> {
    let header = fd.fill_buf()?;
    if header.len() >= 20
        && 0x1f == header[0] && 0x8b == header[1] {
        Ok(FileType::GZip)
   } else if header.len() >= 152
        && b'P' == header[0] && b'K' == header[1]
        && 0x03 == header[2] && 0x04 == header[3] {
        Ok(FileType::Zip)
    } else if header.len() > 257 + 10
        && b'u' == header[257] && b's' == header[258]
        && b't' == header[259] && b'a' == header[260]
        && b'r' == header[261]
        && (
            (0 == header[262] && b'0' == header[263] && b'0' == header[264]) ||
            (b' ' == header[262] && b' ' == header[263] && 0 == header[264])
        ) {
        Ok(FileType::Tar)
    } else if header.len() > 8
        && b'!' == header[0] && b'<' == header[1]
        && b'a' == header[2] && b'r' == header[3]
        && b'c' == header[4] && b'h' == header[5]
        && b'>' == header[6] && b'\n' == header[7] {
        Ok(FileType::Ar)
    } else if header.len() > 40
        && b'B' == header[0] && b'Z' == header[1]
        && b'h' == header[2] && 0x31 == header[3]
        && 0x41 == header[4] && 0x59 == header[5]
        && 0x26 == header[6] {
        Ok(FileType::BZip2)
    } else if header.len() > 6
        && 0xfd == header[0] && b'7' == header[1]
        && b'z' == header[2] && b'X' == header[3]
        && b'Z' == header[4] && 0 == header[5] {
        Ok(FileType::Xz)
    } else if is_probably_tar(header) {
        Ok(FileType::Tar)
    } else {
        Ok(FileType::Other)
    }
}

struct BoxReader<'a> {
    fd: Box<io::BufRead + 'a>
}

impl<'a> io::Read for BoxReader<'a> {
    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
        self.fd.read(&mut buf)
    }
}

fn count_bytes<'a, R: io::Read>(mut fd: &mut R) -> io::Result<u64> {
    let mut block = [0u8; 4096];
    let mut count = 0u64;
    loop {
        let found = fd.read(&mut block[..])?;
        if 0 == found {
            break;
        }
        count += found as u64;
    }
    Ok(count)
}

fn unpack_or_raw<'a, F>(
    fd: Box<io::BufRead + 'a>,
    output: &OutputTo,
    fun: F) -> io::Result<Status>
where F: FnOnce(Box<io::BufRead + 'a>) -> io::Result<Box<io::BufRead + 'a>>
{
    match fun(fd) {
        Ok(inner) => unpack(inner, output),
        Err(e) => {
            output.warn(format!(
                    "thought we could unpack '{:?}' but we couldn't: {}",
                    output.path, e));
            Ok(Status::Rewind)
        }
    }
}

fn unpack<'a>(mut fd: Box<io::BufRead + 'a>, output: &OutputTo) -> io::Result<Status> {
    match identify(&mut fd)? {
        FileType::GZip => {
            unpack_or_raw(fd, output,
                |fd| gzip::Decoder::new(fd).map(
                    |dec| Box::new(io::BufReader::new(dec)) as Box<io::BufRead>))
        },
        FileType::Xz => {
            unpack_or_raw(fd, output,
                |fd| Ok(Box::new(io::BufReader::new(xz2::bufread::XzDecoder::new(fd)))))
        },
        FileType::Ar if output.path.value.ends_with(".deb") => {
            let mut decoder = ar::Archive::new(fd);
            while let Some(entry) = decoder.next_entry() {
                let entry = entry?;
                let new_output = output.with_path(entry.header().identifier());
                unpack(Box::new(io::BufReader::new(entry)), &new_output)?;
            }
            Ok(Status::Done)
        },
        FileType::Tar => {
            let mut decoder = tar::Archive::new(fd);
            for entry in decoder.entries()? {
                let entry = entry?;
                let new_output = output.with_path(entry.path()?.to_str().expect("valid utf-8"));
                unpack(Box::new(io::BufReader::new(entry)), &new_output)?;
            }
            Ok(Status::Done)
        },
        FileType::Zip => {
            let mut temp = tempfile()?;
            io::copy(&mut BoxReader { fd }, &mut temp)?;
            let mut zip = zip::ZipArchive::new(temp)?;

            for i in 0..zip.len() {

                // well, this has gone rather poorly
                let mut new_output = output.with_path(zip.by_index(i)?.name());

                let file_result = {
                    let entry = zip.by_index(i)?;

                    new_output.size = entry.size();
                    new_output.mtime = simple_time_tm(entry.last_modified());
                    let reader = Box::new(io::BufReader::new(entry));
                    unpack(reader, &new_output)?
                };
                match file_result {
                    Status::Done => (),
                    Status::Rewind => {
                        // TODO: update output? ? handling? ???
                        new_output.raw(&mut zip.by_index(i)?)?
                    },
                };
            }

            Ok(Status::Done)
        },

        // TODO: unimplemented!()
        _ => {
            Ok(Status::Rewind)
        },
    }
}

fn output_raw(file: &mut fs::File, output: &OutputTo) -> io::Result<()> {
    file.seek(io::SeekFrom::Start(0))?;
    output.raw(file)?;
    Ok(())
}

fn process_real_path(path: &str) -> io::Result<()> {
    let mut file = fs::File::open(path).unwrap();
    let output = OutputTo::from_file(
        path,
        &file
    )?;

    let attempt = {
        let read = io::BufReader::new(&file);
        unpack(Box::new(read), &output)
    };

    match attempt {
        Ok(Status::Done) => Ok(()),
        Ok(Status::Rewind) => output_raw(&mut file, &output),
        Err(e) => Err(e)
    }
}

fn real_main() -> u8 {

    let matches = App::new("contentin")
                    .arg(Arg::with_name("to-tar")
                         .long("to-tar")
                         .help("emit a tar file")
                         .required(true))
                    .arg(Arg::with_name("INPUT")
                         .required(true)
                         .multiple(true))
                    .get_matches();

    for path in matches.values_of("INPUT").unwrap() {
        if let Err(e) = process_real_path(path) {
            if let Err(_) = writeln!(io::stderr(), "fatal: processing '{}': {}", path, e) {
                return 6;
            }

            return 4;
        }
    }
    return 0;
}

fn main() {
    std::process::exit(real_main() as i32)
}
<|endoftext|>


<commit_before>use std::env;
use std::error;
use std::process;
use std::collections::hash_map::HashMap;

mod add;
mod detail;
mod fat;
mod list;

macro_rules! errorf {
    ($fmt:expr, $($arg:expr),*) => {
        error(From::from(format!(
            $fmt,
            $( $arg ),*
        )));
    }
}

type Command = fn (&[String]) -> Result<(), Box<error::Error>>;
fn get_commands() -> HashMap<&'static str, Command> {
    let mut map = HashMap::new();
    map.insert("add",    add::add_file as Command);
    map.insert("ls",     list::list_files as Command);
    map.insert("detail", detail::detail_file as Command);
    map
}

fn main() {
    let mut arg_iter = env::args();
    let name = arg_iter.next().unwrap();
    let command: String;
    match arg_iter.next() {
        Some(arg) => command = arg,
        None => {
            usage(name);
            process::exit(-1);
        }
    }

    // check for command
    let args: Vec<String> = arg_iter.collect();
    match get_commands().get(&command.as_ref()) {
        Some(cmd_func) =>
            if let Err(e) = cmd_func(&args) {
                error(e);
            },

        None => errorf!("command \"{}\" not recognized", command),
    }
}

fn usage(name: String) {
    println!("{}:", name);
    println!("\tadd <file> <image>: Adds file to image.");
    println!("\tls <image> [flags]: list files on image.");
}

fn error(err: Box<error::Error>) {
    println!("error: {}", err);
    process::exit(-1);
}
<commit_msg>Automate usage generation<commit_after>use std::env;
use std::error;
use std::process;
use std::collections::hash_map::HashMap;

mod add;
mod detail;
mod fat;
mod list;

macro_rules! errorf {
    ($fmt:expr, $($arg:expr),*) => {
        error(From::from(format!(
            $fmt,
            $( $arg ),*
        )));
    }
}

type Command = fn (&[String]) -> Result<(), Box<error::Error>>;
const COMMANDS: &'static [
    (&'static str, &'static str, &'static str, Command)
] = &[
    ("add", "<file> <image>", "Adds file to image.",   add::add_file),
    ("ls",  "<image>", "List files on image.",         list::list_files),
    ("detail",  "<image> <file>", "Show file details", detail::detail_file),
];
fn get_commands() -> HashMap<&'static str, Command> {
    let mut map = HashMap::new();
    map.insert("add",    add::add_file as Command);
    map.insert("ls",     list::list_files as Command);
    map.insert("detail", detail::detail_file as Command);
    map
}

fn main() {
    let mut arg_iter = env::args();
    let name = arg_iter.next().unwrap();
    let command: String;
    match arg_iter.next() {
        Some(arg) => command = arg,
        None => {
            usage(name);
            process::exit(-1);
        }
    }

    // check for command
    let args: Vec<String> = arg_iter.collect();
    match get_commands().get(&command.as_ref()) {
        Some(cmd_func) =>
            if let Err(e) = cmd_func(&args) {
                error(e);
            },

        None => errorf!("command \"{}\" not recognized", command),
    }
}

fn usage(name: String) {
    println!("{}:", name);
    for &(name, usage, description, _) in COMMANDS {
        println!("\t{} {}: {}", name, usage, description);
    }
}

fn error(err: Box<error::Error>) {
    println!("error: {}", err);
    process::exit(-1);
}
<|endoftext|>


<commit_before>#![feature(plugin)]
#![plugin(rocket_codegen)]

extern crate rocket;

extern crate serde_json;
// #[macro_use]
// extern crate serde_derive;

use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use rocket::State;

extern crate hyper;
extern crate hyper_native_tls;

use std::io::Read;

use hyper::Client;
use hyper::net::HttpsConnector;
use hyper_native_tls::NativeTlsClient;

extern crate json;

extern crate time;

extern crate iso_country;
use iso_country::Country;

struct Store(Arc<Mutex<Option<String>>>);

#[get("/")]
fn index(store: State<Store>) -> rocket::response::content::JSON<String> {
    rocket::response::content::JSON(
        store.inner().0
            .lock().unwrap()
            .clone()
            .unwrap_or("{}".to_string())
    )
}

#[inline]
fn process_entry(data: &json::object::Object) -> json::object::Object {
    let mut out = json::object::Object::new();

    for (name, value) in data.iter() {
        if name == "e" {
            out["requests"] = value.clone();
        }

        if name == "n" {
            out["new"] = value.clone();
        }

        if name == "t" {
            out["total"] = value.clone();
        }

        if name == "u" {
            out["unique"] = value.clone();
        }

        if name == "d" {
            out["total_duration"] = value.clone();
        }

        if name == "ds" {
            // ignore
            continue;
        }
    }

    out
}

#[inline]
fn process_ctly(data: &json::JsonValue) -> json::JsonValue {
    let ts = time::now();

    let day = ts.tm_mday.to_string();
    let month = (ts.tm_mon + 1).to_string();
    let year = (ts.tm_year + 1900).to_string();

    let mut data = match data[year][month][day].clone() {
        json::JsonValue::Object(data) => data,
        _ => {
            return json::JsonValue::new_object();
        }
    };

    let mut out = json::object::Object::new();

    let mut per_country = json::object::Object::new();
    let mut intraday_hours = json::object::Object::new();

    for (name, value) in data.iter_mut() {
        let c: Result<i32, _> = name.parse();

        let value: json::object::Object = match *value {
            json::JsonValue::Object(ref mut data) => data.clone(),
            _ => {
                continue;
            }
        };

        if c.is_ok() {
            intraday_hours.insert(&name, json::JsonValue::Object(process_entry(&value)));

            continue;
        }

        // try country
        let c: Result<Country, _> = name.parse();

        if c.is_ok() {
            per_country.insert(&c.unwrap().name(), json::JsonValue::Object(process_entry(&value)));

            continue;
        }

    }

    out["total"] = json::JsonValue::Object(process_entry(&data));

    out["hours"] = json::JsonValue::Object(intraday_hours);
    out["countries"] = json::JsonValue::Object(per_country);

    json::JsonValue::Object(out)
}

fn main() {
    let api_token = match std::env::var("API_TOKEN") {
        Ok(token) => token,
        _ => {
            panic!("Could not find API token!");
        }
    };

    let endpoint_url = format!(
        "https://metrics.psychonautwiki.org/o?api_key={}&app_id=58277520195a624d00fdfaa8&method=users&action=refresh&_=1493844779775",

        api_token
    );

    let data: Option<String> = None;

    // mut
    let root_store = Arc::new(Mutex::new(data));

    /* rocket */

    let store = root_store.clone();

    thread::spawn(move || {
       rocket::ignite()
        .mount("/", routes![index])
        .manage(Store(store))
        .launch();
    });

    /* crawler */

    loop {
        let endpoint_url = endpoint_url.clone();
        let store = root_store.clone();

        let _ = thread::spawn(move || {
            let ssl = NativeTlsClient::new().unwrap();
            let connector = HttpsConnector::new(ssl);
            let client = Client::with_connector(connector);

            let mut res = client.get(&endpoint_url).send().expect("Couldn't send request.");

            let mut buf = String::new();
            res.read_to_string(&mut buf).expect("Couldn't read response.");

            let data = json::parse(&buf).unwrap_or(json::JsonValue::new_object());

            *store.lock().unwrap() = Some(process_ctly(&data).dump());

            thread::sleep(Duration::from_millis(2000));
        }).join();
    }
}<commit_msg>Introduce additional debugging facilities<commit_after>#![feature(plugin)]
#![plugin(rocket_codegen)]

extern crate rocket;

extern crate serde_json;
// #[macro_use]
// extern crate serde_derive;

use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use rocket::State;

extern crate hyper;
extern crate hyper_native_tls;

use std::io::Read;

use hyper::Client;
use hyper::net::HttpsConnector;
use hyper_native_tls::NativeTlsClient;

extern crate json;

extern crate time;

extern crate iso_country;
use iso_country::Country;

struct Store(Arc<Mutex<Option<String>>>);

#[get("/")]
fn index(store: State<Store>) -> rocket::response::content::JSON<String> {
    let shared_data = {
        store
            .inner().0
            .lock()
            .unwrap()
            .clone()
    };

    let data: String = match shared_data {
        Some(data) => data,
        None => {
            println!("Serve error: Using fallback string for None shared data");

            "{}".to_string()
        }
    };

    rocket::response::content::JSON(
        data
    )
}

#[inline]
fn process_entry(data: &json::object::Object) -> json::object::Object {
    let mut out = json::object::Object::new();

    for (name, value) in data.iter() {
        if name == "e" {
            out["requests"] = value.clone();
        }

        if name == "n" {
            out["new"] = value.clone();
        }

        if name == "t" {
            out["total"] = value.clone();
        }

        if name == "u" {
            out["unique"] = value.clone();
        }

        if name == "d" {
            out["total_duration"] = value.clone();
        }

        if name == "ds" {
            // ignore
            continue;
        }
    }

    out
}

#[inline]
fn process_ctly(data: &json::JsonValue) -> json::JsonValue {
    let ts = time::now();

    let day = ts.tm_mday.to_string();
    let month = (ts.tm_mon + 1).to_string();
    let year = (ts.tm_year + 1900).to_string();

    let mut data = match data[year][month][day].clone() {
        json::JsonValue::Object(data) => data,
        _ => {
            return json::JsonValue::new_object();
        }
    };

    let mut out = json::object::Object::new();

    let mut per_country = json::object::Object::new();
    let mut intraday_hours = json::object::Object::new();

    for (name, value) in data.iter_mut() {
        let c: Result<i32, _> = name.parse();

        let value: json::object::Object = match *value {
            json::JsonValue::Object(ref mut data) => data.clone(),
            _ => {
                continue;
            }
        };

        if c.is_ok() {
            intraday_hours.insert(&name, json::JsonValue::Object(process_entry(&value)));

            continue;
        }

        // try country
        let c: Result<Country, _> = name.parse();

        if c.is_ok() {
            per_country.insert(&c.unwrap().name(), json::JsonValue::Object(process_entry(&value)));

            continue;
        }

    }

    out["total"] = json::JsonValue::Object(process_entry(&data));

    out["hours"] = json::JsonValue::Object(intraday_hours);
    out["countries"] = json::JsonValue::Object(per_country);

    json::JsonValue::Object(out)
}

fn main() {
    let api_token = match std::env::var("API_TOKEN") {
        Ok(token) => token,
        _ => {
            panic!("Could not find API token!");
        }
    };

    let endpoint_url = format!(
        "https://metrics.psychonautwiki.org/o?api_key={}&app_id=58277520195a624d00fdfaa8&method=users&action=refresh&_=1493844779775",

        api_token
    );

    let data: Option<String> = None;

    // mut
    let root_store = Arc::new(Mutex::new(data));

    /* rocket */

    let store = root_store.clone();

    thread::spawn(move || {
       rocket::ignite()
        .mount("/", routes![index])
        .manage(Store(store))
        .launch();
    });

    /* crawler */

    loop {
        let endpoint_url = endpoint_url.clone();
        let store = root_store.clone();

        let _ = thread::spawn(move || {
            let ssl = NativeTlsClient::new().unwrap();
            let connector = HttpsConnector::new(ssl);
            let client = Client::with_connector(connector);

            let mut res = client.get(&endpoint_url).send().expect("Connection error: Couldn't send request.");

            let mut buf = String::new();
            res.read_to_string(&mut buf).expect("Connection error: Couldn't read response.");

            let data: json::JsonValue = match json::parse(&buf) {
                Ok(data) => data,
                Err(err) => {
                    panic!("Parse error: {:?}", err);
                }
            };

            *store.lock().unwrap() = Some(process_ctly(&data).dump());

            thread::sleep(Duration::from_millis(2000));
        }).join();
    }
}<|endoftext|>