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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

//! A crate to generate a message factory

use glob::glob;
use regex::Regex;
use std::fs::File;
use std::io::{BufRead, BufReader, Read, Write};
use std::path::PathBuf;

/// protobuf message and file info
#[derive(Debug)]
pub struct ProtoMessageInfo {
    file: PathBuf,
    file_name: String,
    messages: Vec<String>,
}

/// get protobuf message and file info
pub fn get_protos_info(p: &str) -> Vec<ProtoMessageInfo> {
    let mut v = Vec::new();
    let mut path = p.to_string();

    let re = Regex::new(r"message\s+([^\s]+)\s*\{*$").unwrap();

    path.push_str("/*.proto");

    for entry in glob(path.as_str()).expect("Failed to read glob pattern") {
        if let Ok(path) = entry {
            let f = File::open(path.clone()).expect("Failed to open file");
            let reader = BufReader::new(f);
            let mut item = ProtoMessageInfo {
                file: path.clone(),
                file_name: path.file_stem().unwrap().to_str().unwrap().to_string(),
                messages: vec![],
            };

            for line in reader.lines() {
                for caps in re.captures_iter(line.unwrap().as_str()) {
                    item.messages
                        .push(caps.get(1).unwrap().as_str().to_string());
                }
            }
            v.push(item)
        }
    }

    v
}

/// generate factory into `path`
pub fn generate_factory_file(path: &str, v: &Vec<ProtoMessageInfo>) {
    let mut contents = "use std::cell::RefCell;
use std::collections::HashMap;
use protobuf::reflect::MessageDescriptor;
use protobuf::Message;


thread_local! {
    pub static GLOBAL_MAP : RefCell<HashMap<String, &'static MessageDescriptor>> = RefCell::new(HashMap::new());
}

pub fn register_message<M: Message>() {
    GLOBAL_MAP.with(|x| {
        let mut m = x.borrow_mut();
        let name = M::descriptor_static().full_name().to_string();
        if !m.contains_key(&name) {
            m.insert(name, M::descriptor_static());
        }
    })
}

pub fn get_descriptor(full_name: &String) -> Option<&'static MessageDescriptor> {
    GLOBAL_MAP.with(move |x| {
        {
            let m = x.borrow_mut();
            if m.len() == 0 {
                drop(m);
                init_descriptors()
            }
        }
        {
            let m = x.borrow_mut();
            match m.get(full_name) {
                Some(r) => Some(*r),
                None => None,
            }
        }
    })
}".to_string().into_bytes();

    let mut mod_file = File::create((path.to_string() + "/lib.rs").as_str()).unwrap();
    let mut factory_file = File::create((path.to_string() + "/factory.rs").as_str()).unwrap();

    mod_file.write(b"pub mod factory;\n");

    factory_file.write_all(&contents[..]);
    factory_file.write(b"\n\n");

    for item in v.iter() {
        factory_file.write_fmt(format_args!("use crate::{};\n", item.file_name));
        mod_file.write_fmt(format_args!("pub mod {};\n", item.file_name));
    }

    factory_file.write(b"\nfn init_descriptors() {");

    for file in v.iter() {
        for msg in file.messages.iter() {
            factory_file.write_fmt(format_args!(
                "\n    register_message::<{}::{}>();",
                file.file_name, msg
            ));
        }
    }

    factory_file.write(b"\n}\n");
}

/// get proto's filename list
pub fn get_proto_list(v: &Vec<ProtoMessageInfo>) -> Vec<&str> {
    let mut r = Vec::new();

    for f in v.iter() {
        r.push(f.file.to_str().unwrap());
    }

    r
}

//fn main() {
//    let proto_path = "src/";
//
//    let v = get_protos_info(proto_path);
//    let protos = get_proto_list(&v);
//
//    protoc_rust::run(protoc_rust::Args {
//        out_dir: proto_path,
//        input: &protos,
//        includes: &[proto_path],
//        customize: Customize {
//          ..Default::default()
//        },
//    }).expect("protoc");
//
//    generate_factory_file(proto_path, &v);
//}