1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

#![crate_type="dylib"]
#![feature(plugin_registrar, quote, rustc_private)]

extern crate syntax;
extern crate syntax_pos;
extern crate rustc;
extern crate rustc_plugin;
extern crate glob;

use syntax::ast;
use syntax::tokenstream::TokenTree;
use syntax::ext::base::{ExtCtxt, MacResult, DummyResult, MacEager};
use syntax::ext::build::AstBuilder;
use syntax_pos::Span;
use syntax_pos::symbol::Symbol;
use rustc_plugin::Registry;
use syntax::ext::base::*;

use std::fs::{File, metadata};
use std::io::prelude::*;
use std::path::{Path, PathBuf};
use self::glob::glob;
use std::rc::Rc;

fn expand_include_dir(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let dir = match get_single_str_from_tts(cx, sp, tts, "include_dir!") {
        Some(d) => d,
        None => return DummyResult::expr(sp),
    };
    let dir = res_rel_file(cx, sp, Path::new(&dir));

    let expr_new_map = quote_expr!(cx, std::collections::HashMap::new());
    let stmt_let_map = quote_stmt!(cx, let mut map = $expr_new_map;)
        .unwrap();

    let mut stmts = vec![];
    stmts.push(stmt_let_map);

    let mut glob_dir = dir.clone();
    glob_dir.push("**/*");
    let paths = glob(glob_dir.to_str().unwrap()).unwrap();
    for path in paths {
        if let Ok(path) = path {
            let stat = metadata(&path).unwrap();
            if stat.is_file() {
                let mut bytes = Vec::new();
                match File::open(&path).and_then(|mut f| f.read_to_end(&mut bytes)) {
                    Err(e) => {
                        cx.span_err(sp, &format!("couldn't read {}: {}", path.display(), e));
                        return DummyResult::expr(sp);
                    }
                    Ok(..) => {
                        let filename = format!("{}", path.display());
                        cx.codemap().new_filemap_and_lines(&filename, "");

                        let lit_bytes = cx.expr_lit(sp, ast::LitKind::ByteStr(Rc::new(bytes)));

                        let stripped = path.strip_prefix(&dir).unwrap();
                        let lit_path = cx.expr_lit(
                            sp,
                            ast::LitKind::Str(
                                Symbol::intern(stripped.to_str().unwrap()),
                                ast::StrStyle::Cooked,
                            ),
                        );

                        stmts.push(
                            quote_stmt!(cx,
                                        let _ = map.insert(std::path::Path::new($lit_path),
                                                           $lit_bytes.to_vec());).unwrap(),
                        );
                    }
                };
            }
        }
    }

    stmts.push(quote_stmt!(cx, map).unwrap());

    let block = cx.expr_block(cx.block(sp, stmts));

    MacEager::expr(block)
}

// resolve a file-system path to an absolute file-system path (if it
// isn't already)
fn res_rel_file(cx: &mut ExtCtxt, sp: Span, arg: &Path) -> PathBuf {
    // NB: relative paths are resolved relative to the compilation unit
    if !arg.is_absolute() {
        let mut cu = PathBuf::from(&cx.codemap().span_to_filename(sp));
        cu.pop();
        cu.push(arg);
        cu
    } else {
        arg.to_path_buf()
    }
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("include_dir", expand_include_dir);
}