build.rs

 /**
  * Copyright 2024 ByteDance and/or its affiliates
  *
  * Original Files：protoc-gen-ts (https://github.com/thesayyn/protoc-gen-ts)
  * Copyright (c) 2024 Sahin Yort
  * SPDX-License-Identifier: MIT 
 */

use protobuf_codegen::Codegen;
use std::{env, path::Path};
use swc_common::{
    errors::{ColorConfig, Handler},
    sync::Lrc,
    FileName, SourceMap,
};
use swc_ecma_ast::{ClassMember, ModuleItem};
use swc_ecma_parser::{lexer::Lexer, Capturing, Parser, StringInput, Syntax};

struct WktGen {
    types: Vec<(String, String)>,
    contents: String,
}

impl WktGen {
    pub fn new() -> Self {
        WktGen {
            types: vec![],
            contents: String::new(),
        }
    }
    fn ast(&self, input: String) -> Vec<ModuleItem> {
        let cm: Lrc<SourceMap> = Default::default();
        let handler = Handler::with_tty_emitter(ColorConfig::Never, true, true, Some(cm.clone()));
        let fm = cm.new_source_file(FileName::Real("build.ts".into()), input);
        let lexer = Lexer::new(
            Syntax::Typescript(Default::default()),
            Default::default(),
            StringInput::from(&*fm),
            None,
        );

        let capturing = Capturing::new(lexer);
        let mut parser = Parser::new_from(capturing);

        let body = parser
            .parse_typescript_module()
            .map_err(|e| e.into_diagnostic(&handler).emit())
            .expect("Failed to parse module.")
            .body;

        let mut errors = String::new();
        for e in parser.take_errors() {
            errors.push_str(&e.into_diagnostic(&handler).message().to_string());
        }

        if errors.len() > 0 {
            panic!("{}", errors)
        }

        body
    }

    fn json_methods(&self, ast: ModuleItem) -> (String, ClassMember, ClassMember) {
        let decl = ast
            .as_stmt()
            .expect("expected stmt")
            .as_decl()
            .expect("expected decl")
            .as_class()
            .expect("expected class");
        let mut from_json: Option<ClassMember> = None;
        let mut to_json: Option<ClassMember> = None;
        for b in decl.class.body.clone() {
            if b.is_method() {
                let method = b.as_method().unwrap();
                let keyname = method
                    .key
                    .as_ident()
                    .expect("expected ident")
                    .sym
                    .to_string();
                if keyname == "fromJson" {
                    from_json = Some(b);
                } else if keyname == "toJson" {
                    to_json = Some(b);
                }
            }
        }
        if from_json.is_none() || to_json.is_none() {
            panic!("wkt: missing fromJson or toJson");
        }
        return (
            decl.ident.sym.to_string(),
            from_json.unwrap(),
            to_json.unwrap(),
        );
    }

    fn gen_wkt(&mut self, name: String, types: Vec<(String, String, String)>) {
        self.contents.push_str(&format!("pub mod r#{} {{", name));

        for (name, from_json, to_json) in types {
            self.contents.push_str(
                &"
pub mod r#[name] {
    use swc_ecma_ast::ClassMember;
    const FROM_JSON: &str = r#\"from_json_src\"#;
    pub fn from_json() -> ClassMember {
        serde_json::from_str::<ClassMember>(FROM_JSON).unwrap()
    }

    const TO_JSON: &str = r#\"to_json_src\"#;
    pub fn to_json() -> ClassMember {
        serde_json::from_str::<ClassMember>(TO_JSON).unwrap()
    }
}
"
                .replace("[name]", &name)
                .replace("from_json_src", &from_json)
                .replace("to_json_src", &to_json),
            )
        }

        self.contents.push_str("}\n");
    }

    pub fn build(&mut self, name: String, source: String) -> &mut Self {
        let mast = self.ast(source);

        let mut types: Vec<(String, String, String)> = vec![];

        for m in mast {
            let (type_name, from_json, to_json) = self.json_methods(m);
            self.types.push((
                format!("{}.proto#{}", name, type_name),
                format!("r#{}::r#{}", name, type_name.to_lowercase()),
            ));
            types.push((
                type_name.to_lowercase(),
                serde_json::to_string(&from_json).unwrap(),
                serde_json::to_string(&to_json).unwrap(),
            ));
        }

        self.gen_wkt(name, types);
        self
    }

    fn yield_mod(&mut self) {
        let out_dir = env::var_os("OUT_DIR").unwrap();
        let wkt_dir: std::path::PathBuf = Path::new(&out_dir).join("wkt");
        std::fs::create_dir_all(&wkt_dir).expect("failed to create wkt dir");
        let dest_path = wkt_dir.join("mod.rs");
        let mut arms = String::from("");

        for (t1, t2) in self.types.clone() {
            arms.push_str("\t\t");
            arms.push_str(&format!(
                r#""google/protobuf/{}#from_json" => Some({}::from_json()),"#,
                t1, t2
            ));
            arms.push_str("\n");
            arms.push_str("\t\t");
            arms.push_str(&format!(
                r#""google/protobuf/{}#to_json" => Some({}::to_json()),"#,
                t1, t2
            ));
            arms.push_str("\n");
        }
        self.contents.push_str(&format!(
            r#"
pub fn get_member(proto: &str, type_name: &str, member: &str) -> Option<swc_ecma_ast::ClassMember> {{
    match format!("{{}}#{{}}#{{}}", proto, type_name, member).as_str() {{
{}
        _ => None,
    }}
}}
"#,
            arms
        ));
        std::fs::write(&dest_path, &self.contents).unwrap();
    }
}

pub const STRUCT: &str = include_str!("./js/runtime/google_protobuf/struct.ts");
pub const ANY: &str = include_str!("./js/runtime/google_protobuf/any.ts");
pub const WRAPPERS: &str = include_str!("./js/runtime/google_protobuf/wrappers.ts");
// pub const TIMESTAMP: &str = include_str!("./js/runtime/google_protobuf/timestamp.ts");
// pub const DURATION: &str = include_str!("./js/runtime/google_protobuf/duration.ts");
pub const FIELD_MASK: &str = include_str!("./js/runtime/google_protobuf/field_mask.ts");

fn main() {
    Codegen::new()
        .pure()
        .cargo_out_dir("protogen")
        .input("src/descriptor/descriptor.proto")
        .input("src/descriptor/plugin.proto")
        .include("src/descriptor")
        .run_from_script();

        WktGen::new()
        .build("struct".to_string(), STRUCT.to_string())
        .build("any".to_string(), ANY.to_string())
        .build("wrappers".to_string(), WRAPPERS.to_string())
        // .build("timestamp".to_string(), TIMESTAMP.to_string())
        // .build("duration".to_string(), DURATION.to_string())
        .build("field_mask".to_string(), FIELD_MASK.to_string())
        .yield_mod();
}